chore: import upstream snapshot with attribution

.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

.bandit

@@ -0,0 +1,3 @@
+[bandit]
+exclude = tests
+skips = B101,B615,B102,B110

.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

.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

.editorconfig

@@ -0,0 +1,14 @@
+root = true
+
+[*]
+end_of_line = lf
+insert_final_newline = true
+trim_trailing_whitespace = true
+
+[*.py]
+indent_style = space
+indent_size = 4
+
+[**.yml]
+indent_style = space
+indent_size = 2

.gitattributes

@@ -0,0 +1 @@
+data/*.jsonl filter=lfs diff=lfs merge=lfs -text

.github/CODE_OF_CONDUCT.md

@@ -0,0 +1,129 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement on Discord
+at https://discord.gg/QYF8QrtEUm
+
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

.github/CONTRIBUTING.md

@@ -0,0 +1,126 @@
+# Contributing to axolotl
+
+First of all, thank you for your interest in contributing to axolotl! We appreciate the time and effort you're willing to invest in making our project better. This document provides guidelines and information to make the contribution process as smooth as possible.
+
+## Table of Contents
+
+- [Code of Conduct](#code-of-conduct)
+- [Getting Started](#getting-started)
+- [How to Contribute](#how-to-contribute)
+  - [Reporting Bugs](#reporting-bugs)
+  - [Suggesting Enhancements](#suggesting-enhancements)
+  - [Submitting Pull Requests](#submitting-pull-requests)
+- [Style Guidelines](#style-guidelines)
+  - [Code Style](#code-style)
+  - [Commit Messages](#commit-messages)
+- [Additional Resources](#additional-resources)
+
+## 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/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/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
+# 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/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/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
+
+1. Create a new branch for your feature or bugfix. Use a descriptive name like `feature/your-feature-name` or `fix/your-bugfix-name`.
+2. Make your changes, following the [Style Guidelines](#style-guidelines) below.
+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. 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.
+
+#### GPU End-to-End Tests
+
+GPU-heavy CI (the `docker-e2e-tests` and multi-GPU e2e workflows) is opt-in on pull requests: it only runs once a maintainer applies the `run-gpu-tests` label. Labeling starts the GPU suites immediately without re-running the CPU checks, and subsequent pushes to a labeled PR re-run them automatically. Outside of PRs, the `docker-e2e-tests` suite runs on merges to `main`, and the multi-GPU suite runs on its semi-weekly schedule or manual dispatch.
+
+## Style Guidelines
+
+### Code Style
+
+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
+
+Write clear and concise commit messages that briefly describe the changes made in each commit. Use the imperative mood and start with a capitalized verb, e.g., "Add new feature" or "Fix bug in function".
+
+## Additional Resources
+
+- [GitHub Help](https://help.github.com/)
+- [GitHub Pull Request Documentation](https://docs.github.com/en/github/collaborating-with-issues-and-pull-requests)
+- [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!

.github/FUNDING.yml

@@ -0,0 +1,13 @@
+# These are supported funding model platforms
+
+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: # 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: # Replace with up to 4 custom sponsorship URLs e.g., ['link1', 'link2']

.github/ISSUE_TEMPLATE/bug-report.yaml

@@ -0,0 +1,113 @@
+name: Bug Report
+description: File a bug report
+labels: ["bug", "needs triage"]
+body:
+  - type: markdown
+    attributes:
+      value: |
+        ## Before you start
+        Please **make sure you are on the latest version.**
+        If you encountered the issue after you installed, updated, or reloaded, **please try restarting before reporting the bug**.
+
+  - type: checkboxes
+    id: no-duplicate-issues
+    attributes:
+      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/axolotl-ai-cloud/axolotl/labels/bug) didn't find any similar reports."
+          required: true
+
+  - type: textarea
+    id: expected
+    attributes:
+      label: Expected Behavior
+      description: Tell us what **should** happen.
+    validations:
+      required: true
+
+  - type: textarea
+    id: what-happened
+    attributes:
+      label: Current behaviour
+      description: |
+        Tell us what happens instead of the expected behavior.
+        Provide stacktrace and/or screenshots.
+    validations:
+      required: true
+
+  - type: textarea
+    id: reproduce
+    attributes:
+      label: Steps to reproduce
+      description: |
+        Which exact steps can a developer take to reproduce the issue?
+        The more detail you provide, the easier it will be to narrow down and fix the bug.
+        Please paste in tasks and/or queries **as text, not screenshots**.
+      placeholder: |
+        Example of the level of detail needed to reproduce any bugs efficiently and reliably.
+        1. Go to the '...' page.
+        2. Click on the '...' button.
+        3. Scroll down to '...'.
+        4. Observe the error.
+    validations:
+      required: true
+
+  - type: textarea
+    id: config
+    attributes:
+      label: Config yaml
+      description: |
+        Please attach the config yaml!
+      render: yaml
+
+  - type: textarea
+    id: possible-solution
+    attributes:
+      label: Possible solution
+      description: |
+        Not obligatory, but please suggest a fix or reason for the bug, if you have an idea.
+
+
+  - type: checkboxes
+    id: operating-systems
+    attributes:
+      label: Which Operating Systems are you using?
+      description: You may select more than one.
+      options:
+        - label: Linux
+        - label: macOS
+        - label: Windows
+
+  - type: input
+    id: Python-version
+    attributes:
+      label: Python Version
+      description: Which {Programming} version are you using?
+      placeholder: 3.10 / please change accordingly
+    validations:
+      required: true
+
+  - type: input
+    id: axolotl-branch-commit
+    attributes:
+      label: axolotl branch-commit
+      description: On which branch/commit are you?
+      placeholder: main/4d6490b
+    validations:
+      required: true
+
+  - type: checkboxes
+    id: acknowledgements
+    attributes:
+      label: 'Acknowledgements'
+      description: 'Please confirm the following:'
+      options:
+        - label: 'My issue title is concise, descriptive, and in title casing.'
+          required: true
+        - label: 'I have searched the existing issues to make sure this bug has not been reported yet.'
+          required: true
+        - label: 'I am using the latest version of axolotl.'
+          required: true
+        - label: 'I have provided enough information for the maintainers to reproduce and diagnose the issue.'
+          required: true

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1,7 @@
+blank_issues_enabled: false
+contact_links:
+  - name: Ask a question
+    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

.github/ISSUE_TEMPLATE/docs.yml

@@ -0,0 +1,46 @@
+name: Documentation Improvement / Clarity
+description: Make a suggestion to improve the project documentation.
+labels: ['needs triage', 'docs']
+body:
+  - type: markdown
+    attributes:
+      value: '## :book: Documentation :book:'
+  - type: markdown
+    attributes:
+      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/axolotl-ai-cloud/axolotl/issues).
+  - type: textarea
+    attributes:
+      label: What piece of documentation is affected?
+      description: Please link to the article you'd like to see updated.
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: What part(s) of the article would you like to see updated?
+      description: |
+        - Give as much detail as you can to help us understand the change you want to see.
+        - Why should the docs be changed? What use cases does it support?
+        - What is the expected outcome?
+    validations:
+      required: true
+  - type: textarea
+    attributes:
+      label: Additional Information
+      description: Add any other context or screenshots about the feature request here.
+    validations:
+      required: false
+  - type: checkboxes
+    id: acknowledgements
+    attributes:
+      label: 'Acknowledgements'
+      description: 'Please confirm the following:'
+      options:
+        - label: 'My issue title is concise, descriptive, and in title casing.'
+          required: true
+        - label: 'I have searched the existing issues to make sure this feature has not been requested yet.'
+          required: true
+        - label: 'I have provided enough information for the maintainers to understand and evaluate this request.'
+          required: true

.github/ISSUE_TEMPLATE/feature-request.yaml

@@ -0,0 +1,63 @@
+name: Feature Request / Enhancement
+description: Suggest a new feature or feature enhancement for the project
+labels: ["enhancement", "needs triage"]
+body:
+  - type: checkboxes
+    id: no-duplicate-issues
+    attributes:
+      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/axolotl-ai-cloud/axolotl/discussions/categories/ideas) didn't find any similar feature requests."
+          required: true
+        - 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
+    id: feature-description
+    validations:
+      required: true
+    attributes:
+      label: "🔖 Feature description"
+      description: "A clear and concise description of what the feature request is."
+      placeholder: "You should add ..."
+
+  - type: textarea
+    id: solution
+    validations:
+      required: true
+    attributes:
+      label: "✔️ Solution"
+      description: "A clear and concise description of what you want to happen, and why."
+      placeholder: "In my use-case, ..."
+
+  - type: textarea
+    id: alternatives
+    validations:
+      required: false
+    attributes:
+      label: "❓ Alternatives"
+      description: "A clear and concise description of any alternative solutions or features you've considered."
+      placeholder: "I have considered ..."
+
+  - type: textarea
+    id: additional-context
+    validations:
+      required: false
+    attributes:
+      label: "📝 Additional Context"
+      description: "Add any other context or screenshots about the feature request here."
+      placeholder: "..."
+
+  - type: checkboxes
+    id: acknowledgements
+    attributes:
+      label: 'Acknowledgements'
+      description: 'Please confirm the following:'
+      options:
+        - label: 'My issue title is concise, descriptive, and in title casing.'
+          required: true
+        - label: 'I have searched the existing issues to make sure this feature has not been requested yet.'
+          required: true
+        - label: 'I have provided enough information for the maintainers to understand and evaluate this request.'
+          required: true

.github/PULL_REQUEST_TEMPLATE.md

@@ -0,0 +1,32 @@
+<!--- Provide a general summary of your changes in the Title above -->
+
+# Description
+
+<!--- Describe your changes in detail -->
+
+## Motivation and Context
+
+<!--- Why is this change required? What problem does it solve? -->
+<!--- If it fixes an open issue, please link to the issue here. -->
+
+## How has this been tested?
+
+<!--- Please describe in detail how you tested your changes. -->
+<!--- 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
+
+<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
+
+## Social Handles (Optional)
+
+<!-- Thanks for submitting a bugfix or enhancement. -->
+<!-- We'd love to show our thanks to you on Twitter & Discord if you provide your handle -->

.github/SECURITY.md

@@ -0,0 +1,9 @@
+# Security Policy
+
+## Supported Versions
+
+Due to the nature of the fast development that is happening in this project, only the latest released version can be supported.
+
+## Reporting a Vulnerability
+
+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.

.github/SUPPORT.md

@@ -0,0 +1,10 @@
+# Support
+
+If you need help with this project or have questions, please:
+
+1. Check the documentation.
+2. Search the existing issues and pull requests.
+3. Create a new issue if your question is not answered or your problem is not solved.
+4. Have a look in the [Discord server](https://discord.gg/HhrNrHJPRb)
+
+Please note that this project is maintained by volunteers who have limited availability. We'll do our best to address your questions and concerns in a timely manner.

.github/dependabot.yml

@@ -0,0 +1,8 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "weekly"
+    cooldown:
+      default-days: 7

.github/release-drafter.yml

@@ -0,0 +1,31 @@
+name-template: 'v$RESOLVED_VERSION'
+tag-template: 'v$RESOLVED_VERSION'
+categories:
+  - title: '🚀 Features'
+    labels:
+      - 'feature'
+      - 'enhancement'
+  - title: '🐛 Bug Fixes'
+    labels:
+      - 'fix'
+      - 'bugfix'
+      - 'bug'
+  - title: '🧰 Maintenance'
+    label: 'chore'
+change-template: '- $TITLE @$AUTHOR (#$NUMBER)'
+change-title-escapes: '\<*_&' # You can add # and @ to disable mentions, and add ` to disable code blocks.
+version-resolver:
+  major:
+    labels:
+      - 'major'
+  minor:
+    labels:
+      - 'minor'
+  patch:
+    labels:
+      - 'patch'
+  default: patch
+template: |
+  ## What’s Changed
+
+  $CHANGES

.github/workflows/base.yml

@@ -0,0 +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-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: "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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
+        with:
+          images: |
+            axolotlai/axolotl-base-uv
+      - name: Login to Docker Hub
+        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@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
+      - name: Build
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
+        with:
+          context: .
+          file: ./docker/${{ matrix.dockerfile }}
+          platforms: ${{ matrix.platforms }}
+          push: ${{ github.event_name != 'pull_request' }}
+          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 }}

.github/workflows/docker-e2e.yml

@@ -0,0 +1,279 @@
+name: docker-e2e-tests
+on:
+  # GPU e2e runs unconditionally on push/merge to main and in the merge queue.
+  # On PRs the jobs are gated behind the `run-gpu-tests` label (see
+  # gate-skip-e2e below); the `labeled` trigger picks the PR up as soon as a
+  # maintainer applies the label. These jobs live in their own workflow so that
+  # labeling doesn't re-run the CPU jobs in tests.yml.
+  merge_group:
+  push:
+    branches:
+      - "main"
+    paths:
+      - "**.py"
+      - "pyproject.toml"
+      - ".github/workflows/*.yml"
+      - "cicd/cicd.sh"
+      - "cicd/cicd_cuda_kernels.sh"
+      - "cicd/Dockerfile-uv.jinja"
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review, labeled]
+    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. `labeled` events for
+# unrelated labels get their own no-op group so they can't cancel a live run.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ (github.event.action == 'labeled' && github.event.label.name != 'run-gpu-tests') && 'label-noop' || 'e2e' }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+
+env:
+  TRANSFORMERS_IS_CI: "yes"
+  UV_SYSTEM_PYTHON: "1"
+
+jobs:
+  gate-skip-e2e:
+    # Gates the whole matrix: on PRs, require the `run-gpu-tests` label and skip
+    # drafts. A `labeled` event for any other label is a no-op so it doesn't
+    # restart a matrix that is already in flight.
+    if: >
+      github.event_name != 'pull_request' ||
+      (
+        !github.event.pull_request.draft &&
+        contains(github.event.pull_request.labels.*.name, 'run-gpu-tests') &&
+        (github.event.action != 'labeled' || github.event.label.name == 'run-gpu-tests')
+      )
+    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 || '';
+            }
+            // don't scan the PR body: bots (e.g. CodeRabbit summaries) quote
+            // the token when describing CI, which false-positives the skip
+            const title = context.payload.pull_request?.title || '';
+            const skip = token.test(msg) || token.test(title);
+            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: [gate-skip-e2e]
+
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.1
+            num_gpus: 1
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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 == '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
+    # 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: [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:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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: [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.1
+            num_gpus: 1
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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.1
+            num_gpus: 1
+            axolotl_extras:
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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
+      - name: Run tests job on Modal
+        run: |
+          modal run -m cicd.cleanup

.github/workflows/docs.yml

@@ -0,0 +1,49 @@
+name: Publish Docs
+on:
+  push:
+    branches:
+      - main
+    # src/axolotl is included because quartodoc renders API docs from the
+    # whole package; without it the published API reference goes stale (or
+    # silently breaks) when source modules change
+    paths:
+      - '**/*.md'
+      - '**/*.qmd'
+      - '_quarto.yml'
+      - 'docs/**'
+      - 'src/axolotl/**.py'
+      - '.github/workflows/docs.yml'
+
+permissions:
+    contents: write
+    pages: write
+
+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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+          with:
+            persist-credentials: false
+        - 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: Publish to GitHub Pages (and render)
+          uses: quarto-dev/quarto-actions/publish@8a96df13519ee81fd526f2dfca5962811136661b # v2.2.0
+          with:
+            target: gh-pages
+          env:
+            GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

.github/workflows/lint.yml

@@ -0,0 +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:
+       - "*.[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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          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: check-cli-config-options

.github/workflows/main.yml

@@ -0,0 +1,210 @@
+name: ci-cd
+
+on:
+  push:
+    branches:
+      - "main"
+    tags:
+      - "v*"
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+permissions:
+  contents: read
+
+jobs:
+  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: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
+            axolotl_extras:
+            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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
+        with:
+          images: |
+            axolotlai/axolotl-uv
+          tags: |
+            type=ref,event=branch
+            type=pep440,pattern={{version}}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
+      - name: Login to Docker Hub
+        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@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
+        with:
+          context: .
+          platforms: ${{ matrix.platforms }}
+          build-args: |
+            BASE_TAG=${{ github.ref_type == 'tag' && 'main' || github.ref_name }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}
+            CUDA=${{ matrix.cuda }}
+            PYTORCH_VERSION=${{ matrix.pytorch }}
+            AXOLOTL_ARGS=${{ matrix.axolotl_args }}
+            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: |
+            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-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:
+            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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.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:
+            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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
+        with:
+          images: |
+            axolotlai/axolotl-cloud-term
+          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@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
+      - name: Build
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
+        with:
+          context: .
+          platforms: linux/amd64,linux/arm64
+          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-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: |
+             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 }}

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

@@ -0,0 +1,83 @@
+name: docker-multigpu-tests-biweekly
+
+on:
+  pull_request:
+    # on PRs the job is gated behind the `run-gpu-tests` label; `labeled`
+    # picks the PR up when a maintainer applies it. schedule/dispatch runs
+    # are not gated.
+    types: [opened, synchronize, reopened, ready_for_review, labeled]
+    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. `labeled` events for
+# unrelated labels get their own no-op group so they can't cancel a live run.
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}-${{ (github.event.action == 'labeled' && github.event.label.name != 'run-gpu-tests') && 'label-noop' || 'e2e' }}
+  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}
+
+permissions:
+  contents: read
+
+env:
+  MODAL_IMAGE_BUILDER_VERSION: "2025.06"
+
+jobs:
+  test-axolotl-multigpu:
+    if: >
+      github.repository_owner == 'axolotl-ai-cloud' &&
+      (
+        github.event_name != 'pull_request' ||
+        (
+          !github.event.pull_request.draft &&
+          contains(github.event.pull_request.labels.*.name, 'run-gpu-tests') &&
+          (github.event.action != 'labeled' || github.event.label.name == 'run-gpu-tests')
+        )
+      )
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.12.1
+            axolotl_extras:
+            # axolotl_extras: fbgemm-gpu
+            num_gpus: 2
+    runs-on: [self-hosted, modal]
+    timeout-minutes: 120
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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

.github/workflows/nightlies.yml

@@ -0,0 +1,120 @@
+name: docker-nightlies
+
+on:
+  workflow_dispatch:
+  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.head_commit.message, '[skip docker]') && github.repository_owner == 'axolotl-ai-cloud' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - cuda: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
+            axolotl_extras:
+          - 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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
+        with:
+          images: |
+            axolotlai/axolotl
+          tags: |
+            type=raw,value={{ branch }}-{{ date 'YYYYMMDD' }}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
+      - name: Login to Docker Hub
+        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@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5.4.0
+        with:
+          context: .
+          build-args: |
+            BASE_TAG=${{ 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-uv
+          push: ${{ github.event_name != 'pull_request' }}
+          tags: |
+            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.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: 130
+            cuda_version: 13.0.0
+            python_version: "3.12"
+            pytorch: 2.11.0
+            axolotl_extras:
+          - 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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5.10.0
+        with:
+          images: |
+            axolotlai/axolotl-cloud
+          tags: |
+            type=raw,value={{ branch }}-{{ date 'YYYYMMDD' }}
+      - 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@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0
+      - name: Build
+        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-uv
+          push: ${{ github.event_name != 'pull_request' }}
+          tags: |
+             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 }}

.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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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.

.github/workflows/preview-docs.yml

@@ -0,0 +1,87 @@
+name: Preview
+on:
+  workflow_dispatch:
+  pull_request:
+    types: [opened, synchronize, reopened, ready_for_review]
+
+    # Run the workflow only when one of these files changes. src/axolotl is
+    # included because quartodoc renders API docs from the whole package —
+    # removing a module it tracks breaks the build (see _quarto.yml sections)
+    paths:
+      - '**/*.md'      # any Markdown file
+      - '**/*.qmd'     # any Quarto file
+      - '_quarto.yml'
+      - 'docs/**'
+      - 'src/axolotl/**.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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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
+
+      # everything above runs credential-free on every PR (including forks) so
+      # a broken docs build fails CI; publishing needs the Netlify secrets,
+      # which only same-repo branches get
+      - name: Netlify Publish
+        uses: nwtgck/actions-netlify@d22a32a27c918fe470bbc562e984f80ec48c2668 # v4.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 }}

.github/workflows/pypi.yml

@@ -0,0 +1,70 @@
+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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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:
+      contents: read
+      id-token: write # IMPORTANT: this permission is mandatory for trusted publishing
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+
+      - name: Setup Python
+        uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          python-version: "3.11"
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@37802adc94f370d6bfd71619e3f0bf239e1f3b78 # v7.6.0
+        with:
+          enable-cache: false
+
+      - name: Extract tag name
+        id: tag
+        run: echo "TAG_NAME=$(echo $GITHUB_REF | cut -d / -f 3)" >> "$GITHUB_OUTPUT"
+
+      - name: Update version in VERSION file
+        run: |
+          echo "${{ steps.tag.outputs.TAG_NAME }}" | sed 's/^v//' > VERSION
+
+      - 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@cef221092ed1bacb1cc03d23a2d87d1d172e277b # v1.14.0

.github/workflows/tests-nightly.yml

@@ -0,0 +1,290 @@
+name: Tests Nightly against upstream main
+on:
+  # nightly-only: runs against upstream HF main, so PR triggers just burn GPU
+  # on failures the PR can't cause. Use workflow_dispatch to validate changes
+  # to this file or the cicd/ harness.
+  workflow_dispatch:
+  schedule:
+    - cron: "0 0 * * *" # Runs at 00:00 UTC every day
+
+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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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"]
+        pytorch_version: ["2.11.0", "2.12.1", "2.13.0"]
+    timeout-minutes: 20
+
+    steps:
+      - name: Check out repository code
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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.1
+            num_gpus: 1
+            axolotl_extras:
+            nightly_build: "true"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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.1
+            num_gpus: 2
+            axolotl_extras:
+            nightly_build: "true"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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.1
+            num_gpus: 1
+            axolotl_extras:
+            nightly_build: "true"
+    steps:
+      - name: Checkout
+        uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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

.github/workflows/tests.yml

@@ -0,0 +1,264 @@
+name: Tests
+on:
+  # check on push/merge to main, PRs, and manual triggers
+  merge_group:
+  push:
+    branches:
+      - "main"
+    paths:
+      - "**.py"
+      - "pyproject.toml"
+      - ".github/workflows/*.yml"
+  pull_request:
+    # no `labeled` here: the label-gated GPU jobs live in docker-e2e.yml so
+    # labeling a PR doesn't re-run the CPU jobs
+    types: [opened, synchronize, reopened, ready_for_review]
+    paths:
+      - "**.py"
+      - "pyproject.toml"
+      - ".github/workflows/*.yml"
+  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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0
+        with:
+          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"]
+        pytorch_version: ["2.11.0", "2.12.1", "2.13.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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0
+        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.12"]
+        pytorch_version: ["2.11.0", "2.12.1", "2.13.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@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
+        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 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 -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

.gitignore

@@ -0,0 +1,196 @@
+**/axolotl.egg-info
+configs
+last_run_prepared/
+outputs
+.vscode
+_site/
+
+# 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
+uv.lock
+
+# 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/
+venv3.10/
+
+# 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/
+
+# WandB
+# wandb creates a folder to store logs for training runs
+wandb
+
+# Runs
+lora-out/*
+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

.mypy.ini

@@ -0,0 +1,58 @@
+[mypy]
+plugins = pydantic.mypy
+exclude = venv
+
+[mypy-alpaca_lora_4bit.*]
+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
+
+[mypy-flash_attn.*]
+ignore_missing_imports = True
+
+[mypy-huggingface_hub]
+ignore_missing_imports = True
+
+[mypy-transformers.*]
+ignore_missing_imports = True
+
+[mypy-peft]
+ignore_missing_imports = True
+
+[mypy-wandb]
+ignore_missing_imports = True
+
+[mypy-bitsandbytes]
+ignore_missing_imports = True
+
+[mypy-requests]
+ignore_missing_imports = True
+
+[mypy-datasets]
+ignore_missing_imports = True
+
+[mypy-fire]
+ignore_missing_imports = True
+
+[mypy-setuptools]
+ignore_missing_imports = True
+
+[mypy-addict]
+ignore_missing_imports = True
+
+[mypy-xformers.*]
+ignore_missing_imports = True

.pre-commit-config.yaml

@@ -0,0 +1,43 @@
+default_language_version:
+    python: python3
+
+repos:
+-   repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+    -   id: check-yaml
+    -   id: end-of-file-fixer
+    -   id: trailing-whitespace
+    -   id: no-commit-to-branch
+        args: ['--branch', 'main']
+-   repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.15.8
+    hooks:
+    -   id: ruff
+        args: [--fix]
+    -   id: ruff-format
+-   repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.19.1
+    hooks:
+    - id: mypy
+      additional_dependencies:
+        [
+            'types-PyYAML',
+            'pydantic>=2.5.3',
+        ]
+-   repo: https://github.com/PyCQA/bandit
+    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

.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

.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"]

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

.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
+        }
+      }
+    ]
+  }
+}

.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

.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}

.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})

.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|>"
+      }
+    }
+  }
+}

.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()

.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)

.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"
+    ]
+  }
+}

.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"
+    ]
+  }
+}

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

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"

CLAUDE.md

@@ -0,0 +1 @@
+AGENTS.md

CNAME

@@ -0,0 +1 @@
+docs.axolotl.ai

FAQS.md

@@ -0,0 +1,7 @@
+# FAQs
+
+- Can you train StableLM with this? Yes, but only with a single GPU atm. Multi GPU support is coming soon! Just waiting on this [PR](https://github.com/huggingface/transformers/pull/22874)
+- Will this work with Deepspeed? That's still a WIP, but setting `export ACCELERATE_USE_DEEPSPEED=true` should work in some cases
+- `Error invalid argument at line 359 in file /workspace/bitsandbytes/csrc/pythonInterface.c`
+`/arrow/cpp/src/arrow/filesystem/s3fs.cc:2598:  arrow::fs::FinalizeS3 was not called even though S3 was initialized.`
+This could lead to a segmentation fault at exit. Try reinstalling bitsandbytes and transformers from source.

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.

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

README.md

@@ -0,0 +1,239 @@
+<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/docker-e2e.yml/badge.svg" alt="docker-e2e-tests">
+    <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>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>
+
+## ✨ Overview
+
+Axolotl is a free and open-source tool designed to streamline post-training and fine-tuning for the latest large language models (LLMs).
+
+Features:
+
+- **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.
+
+
+
+## 🚀 Quick Start - LLM Fine-tuning in Minutes
+
+**Requirements**:
+
+- NVIDIA GPU (Ampere or newer for `bf16` and Flash Attention) or AMD GPU
+- Python >=3.11 (3.12 recommended)
+- PyTorch ≥2.11.0
+
+### Google Colab
+
+[![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)
+
+### Installation
+
+```bash
+# install uv if you don't already have it installed (restart shell after)
+curl -LsSf https://astral.sh/uv/install.sh | sh
+
+# change depending on system
+export UV_TORCH_BACKEND=cu130
+
+# create a new virtual environment
+uv venv --python 3.12
+source .venv/bin/activate
+
+uv pip install torch==2.12.0 torchvision
+uv pip install --no-build-isolation axolotl[deepspeed]
+
+# Download example axolotl configs, deepspeed configs
+axolotl fetch examples
+axolotl fetch deepspeed_configs  # OPTIONAL
+```
+
+#### Using Docker
+
+Installing with Docker can be less error prone than installing in your own environment.
+```bash
+docker run --gpus '"all"' --ipc=host --rm -it axolotlai/axolotl:main-latest
+```
+
+Other installation approaches are described [here](https://docs.axolotl.ai/docs/installation.html).
+
+#### Cloud Providers
+
+<details>
+
+- [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)
+
+</details>
+
+### Your First Fine-tune
+
+```bash
+# Fetch axolotl examples
+axolotl fetch examples
+
+# Or, specify a custom path
+axolotl fetch examples --dest path/to/folder
+
+# Train a model using LoRA
+axolotl train examples/llama-3/lora-1b.yml
+```
+
+That's it! Check out our [Getting Started Guide](https://docs.axolotl.ai/docs/getting-started.html) for a more detailed walkthrough.
+
+
+## 📚 Documentation
+
+- [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
+
+## AI Agent Support
+
+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.
+
+```bash
+# Show overview and available training methods
+axolotl agent-docs
+
+# 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
+
+# Dump config schema for programmatic use
+axolotl config-schema
+axolotl config-schema --field adapter
+```
+
+If you're working with the source repo, agent docs are also available at `docs/agents/` and the project overview is in `AGENTS.md`.
+
+## 🤝 Getting Help
+
+- 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
+
+## 🌟 Contributing
+
+Contributions are welcome! Please see our [Contributing Guide](https://github.com/axolotl-ai-cloud/axolotl/blob/main/.github/CONTRIBUTING.md) for details.
+
+## 📈 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).
+
+## ❤️ Sponsors
+
+Interested in sponsoring? Contact us at [wing@axolotl.ai](mailto:wing@axolotl.ai)
+
+## 📝 Citing Axolotl
+
+If you use Axolotl in your research or projects, please cite it as follows:
+
+```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}
+}
+```
+
+## 📜 License
+
+This project is licensed under the Apache 2.0 License - see the [LICENSE](LICENSE) file for details.

README.wehub.md

@@ -0,0 +1,7 @@
+# WeHub 来源说明
+
+- 原始项目：`axolotl-ai-cloud/axolotl`
+- 原始仓库：https://github.com/axolotl-ai-cloud/axolotl
+- 导入方式：上游默认分支的最新快照
+- 原作者、版权和许可证信息以原始仓库及本仓库 LICENSE 为准
+- 本文件仅用于记录来源，不代表 WeHub 是原项目作者

VERSION

@@ -0,0 +1 @@
+0.17.0.dev

_quarto.yml

@@ -0,0 +1,569 @@
+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.qwen3_5
+        - integrations.liger.models.qwen3_5_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: "We make fine-tuning accessible, scalable, and fun"
+  favicon: favicon.jpg
+
+  google-analytics: "G-9KYCVJBNMQ"
+
+  navbar:
+    logo: image/axolotl_logo_digital_white.svg
+    title: false
+    background: dark
+    pinned: false
+    collapse: false
+    tools:
+    - icon: twitter
+      href: https://twitter.com/axolotl_ai
+    - icon: github
+      href: https://github.com/axolotl-ai-cloud/axolotl/
+    - icon: discord
+      href: https://discord.gg/7m9sfhzaf3
+
+  sidebar:
+      pinned: true
+      collapse-level: 2
+      style: docked
+      contents:
+        - text: Home
+          href: index.qmd
+
+        - section: "Getting Started"
+          contents:
+            - 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: "Deployments"
+          contents:
+            - 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: 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

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

cicd/cicd.sh

@@ -0,0 +1,95 @@
+#!/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
+# 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

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

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()

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 {} \;

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()

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()

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()

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

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}")

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

deepspeed_configs/zero1.json

@@ -0,0 +1,23 @@
+{
+  "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
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}

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
+}

deepspeed_configs/zero2.json

@@ -0,0 +1,27 @@
+{
+  "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
+}

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
+}

deepspeed_configs/zero3.json

@@ -0,0 +1,31 @@
+{
+  "zero_optimization": {
+    "stage": 3,
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "sub_group_size": 0,
+    "reduce_bucket_size": "auto",
+    "stage3_prefetch_bucket_size": "auto",
+    "stage3_param_persistence_threshold": "auto",
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": 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
+}

deepspeed_configs/zero3_bf16.json

@@ -0,0 +1,22 @@
+{
+  "zero_optimization": {
+    "stage": 3,
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "sub_group_size": 0,
+    "reduce_bucket_size": "auto",
+    "stage3_prefetch_bucket_size": "auto",
+    "stage3_param_persistence_threshold": "auto",
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
+  },
+  "bf16": {
+    "enabled": true
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}

deepspeed_configs/zero3_bf16_cpuoffload_all.json

@@ -0,0 +1,32 @@
+{
+  "zero_force_ds_cpu_optimizer": false,
+  "zero_allow_untested_optimizer": true,
+  "zero_optimization": {
+    "stage": 3,
+    "offload_optimizer": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "sub_group_size": 0,
+    "reduce_bucket_size": "auto",
+    "stage3_prefetch_bucket_size": "auto",
+    "stage3_param_persistence_threshold": "auto",
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
+  },
+  "bf16": {
+    "enabled": true
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}

deepspeed_configs/zero3_bf16_cpuoffload_params.json

@@ -0,0 +1,28 @@
+{
+  "zero_force_ds_cpu_optimizer": false,
+  "zero_allow_untested_optimizer": true,
+  "zero_optimization": {
+    "stage": 3,
+    "offload_param": {
+      "device": "cpu",
+      "pin_memory": true
+    },
+    "overlap_comm": true,
+    "contiguous_gradients": true,
+    "sub_group_size": 0,
+    "reduce_bucket_size": "auto",
+    "stage3_prefetch_bucket_size": "auto",
+    "stage3_param_persistence_threshold": "auto",
+    "max_live_parameters": 0,
+    "max_reuse_distance": 0,
+    "gather_16bit_weights_on_model_save": true
+  },
+  "bf16": {
+    "enabled": true
+  },
+  "gradient_accumulation_steps": "auto",
+  "gradient_clipping": "auto",
+  "train_batch_size": "auto",
+  "train_micro_batch_size_per_gpu": "auto",
+  "wall_clock_breakdown": false
+}

devtools/README.md

@@ -0,0 +1 @@
+This directory contains example config files that might be useful for debugging. Please see [docs/debugging.qmd](../docs/debugging.qmd) for more information.

devtools/dev_chat_template.yml

@@ -0,0 +1,48 @@
+# Example config for debugging the chat_template prompt format
+base_model: TinyLlama/TinyLlama-1.1B-Chat-v1.0
+model_type: LlamaForCausalLM
+tokenizer_type: LlamaTokenizer
+
+load_in_8bit: true
+load_in_4bit: false
+
+datasets:
+  - 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_num_proc: 1
+
+sequence_len: 4096
+sample_packing: false
+pad_to_sequence_len: true
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_fan_in_fan_out:
+
+micro_batch_size: 1
+num_epochs: 1
+max_steps: 10
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: false
+fp16: true
+tf32: false
+
+gradient_checkpointing: true
+logging_steps: 1
+flash_attention: true
+
+warmup_steps: 10
+weight_decay: 0.0

docker-compose.yaml

@@ -0,0 +1,25 @@
+# version: '3.8'
+services:
+  axolotl:
+    build:
+      context: .
+      dockerfile: ./docker/Dockerfile-uv
+    volumes:
+      - .:/workspace/axolotl
+      - ~/.cache/huggingface/:/root/.cache/huggingface/
+    # set environment variables
+    environment:
+      # Set environment variables
+      - GIT_AUTHOR_NAME=${GIT_AUTHOR_NAME}
+      - GIT_AUTHOR_EMAIL=${GIT_AUTHOR_EMAIL}
+      - GIT_COMMITTER_NAME=${GIT_COMMITTER_NAME}
+      - GIT_COMMITTER_EMAIL=${GIT_COMMITTER_EMAIL}
+      - WANDB_API_KEY=${WANDB_API_KEY}
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              # count: 1
+              capabilities: [gpu]
+    command: tail -f /dev/null

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"]

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"]

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"]

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

docs/.gitignore

@@ -0,0 +1,7 @@
+/.quarto/
+_site/
+/api/*.qmd
+/api/*.html
+config-reference.qmd
+models/**/*.qmd
+models/**/*.html

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.

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

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

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 |

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

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

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

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

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. 🚂

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.
+
+:::

docs/batch_vs_grad.qmd

@@ -0,0 +1,59 @@
+---
+title: Batch size vs Gradient accumulation
+description: Understanding of batch size and gradient accumulation steps
+---
+
+Gradient accumulation means accumulating gradients over several mini-batches and updating the model weights afterward. When the samples in each batch are diverse, this technique doesn't significantly impact learning.
+
+This method allows for effective training with larger effective batch sizes without needing proportionally larger memory. Here's why:
+
+1. **Memory Consumption with Batch Size**: The primary reason increasing the batch size impacts memory is due to the storage requirements for intermediate activations. When you forward propagate a batch through a network, you have to store the activations at each layer for each sample in the batch, because these activations are used during backpropagation to compute gradients. Therefore, larger batches mean more activations, leading to greater GPU memory consumption.
+
+2. **Gradient Accumulation**: With gradient accumulation, you're effectively simulating a larger batch size by accumulating gradients over several smaller batches (or micro-batches). However, at any given time, you're only forward and backward propagating a micro-batch. This means you only store activations for the micro-batch, not the full accumulated batch. As a result, you can simulate the effect of a larger batch size without the memory cost of storing activations for a large batch.
+
+**Example 1:**
+Micro batch size: 3
+Gradient accumulation steps: 2
+Number of GPUs: 3
+Total batch size = 3 * 2 * 3 = 18
+
+```
+| GPU 1          | GPU 2          | GPU 3          |
+|----------------|----------------|----------------|
+| S1, S2, S3     | S4, S5, S6     | S7, S8, S9     |
+| e1, e2, e3     | e4, e5, e6     | e7, e8, e9     |
+|----------------|----------------|----------------|
+| → (accumulate) | → (accumulate) | → (accumulate) |
+|----------------|----------------|----------------|
+| S10, S11, S12  | S13, S14, S15  | S16, S17, S18  |
+| e10, e11, e12  | e13, e14, e15  | e16, e17, e18  |
+|----------------|----------------|----------------|
+| → (apply)      | → (apply)      | → (apply)      |
+
+Accumulated gradient for the weight w1 after the second iteration (considering all GPUs):
+Total gradient for w1 = e1 + e2 + e3 + e4 + e5 + e6 + e7 + e8 + e9 + e10 + e11 + e12 + e13 + e14 + e15 + e16 + e17 + e18
+
+Weight update for w1:
+w1_new = w1_old - learning rate x (Total gradient for w1 / 18)
+```
+
+**Example 2:**
+Micro batch size: 2
+Gradient accumulation steps: 1
+Number of GPUs: 3
+Total batch size = 2 * 1 * 3 = 6
+
+```
+| GPU 1     | GPU 2     | GPU 3     |
+|-----------|-----------|-----------|
+| S1, S2    | S3, S4    | S5, S6    |
+| e1, e2    | e3, e4    | e5, e6    |
+|-----------|-----------|-----------|
+| → (apply) | → (apply) | → (apply) |
+
+Accumulated gradient for the weight w1 (considering all GPUs):
+Total gradient for w1 = e1 + e2 + e3 + e4 + e5 + e6
+
+Weight update for w1:
+w1_new = w1_old - learning rate × (Total gradient for w1 / 6)
+```

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).

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

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

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)
+
+:::

docs/dataset-formats/conversation.qmd

@@ -0,0 +1,452 @@
+---
+title: Conversation
+description: Conversation format for supervised fine-tuning.
+order: 3
+---
+
+## chat_template
+
+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"}
+{"messages": [{"role": "...", "content": "..."}, {"role": "...", "content": "..."}, ...]}
+```
+
+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: ...
+    type: sharegpt
+    conversation: chatml
+
+# new (if using tokenizer's chat_template)
+datasets:
+  - path: ...
+    type: chat_template
+
+    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
+```
+
+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": [
+    {"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}
+  ]
+}
+```
+
+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
+```
+
+::: {.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"}
+{
+  "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}
+      ]
+    }
+  ]
+}
+```
+
+The configuration is the same as standard `chat_template` — no extra fields needed:
+
+```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"}
+{
+  "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}
+      ]
+    }
+  ]
+}
+```
+
+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.
+:::
+
+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": [{"role": "...", "value": "..."}]}
+```