feat: add baseten to lmeval

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

@@ -44,7 +44,7 @@ jobs:
             cuda_version: 12.8.1
             python_version: "3.11"
             pytorch: 2.8.0
-            axolotl_extras: fbgemm-gpu
+            axolotl_extras:
             num_gpus: 2
             nightly_build: "true"
     runs-on: [self-hosted, modal]

.github/workflows/tests.yml

@@ -303,8 +303,7 @@ jobs:
             python_version: "3.11"
             pytorch: 2.8.0
             num_gpus: 1
-            gpu_type: "B200"
-            axolotl_extras: fbgemm-gpu
+            axolotl_extras:
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -325,7 +324,6 @@ jobs:
           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 "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
           echo "E2E_DOCKERFILE=${{ matrix.dockerfile || 'Dockerfile.jinja'}}" >> $GITHUB_ENV
       - name: Run tests job on Modal

.gitignore

@@ -190,6 +190,3 @@ out/
 
 # vim
 *.swp
-
-# scm auto-versioning
-src/axolotl/_version.py

.pre-commit-config.yaml

@@ -11,7 +11,7 @@ repos:
     -   id: no-commit-to-branch
         args: ['--branch', 'main']
 -   repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.12.12
+    rev: v0.12.9
     hooks:
     -   id: ruff
         args: [--fix]

CITATION.cff

@@ -1,6 +1,6 @@
 cff-version: 1.2.0
 type: software
-title: "Axolotl: Open Source LLM Post-Training"
+title: "Axolotl: Post-Training for AI Models"
 message: "If you use this software, please cite it as below."
 authors:
   - name: "Axolotl maintainers and contributors"

README.md

@@ -5,9 +5,6 @@
         <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">
@@ -20,7 +17,6 @@
     <br/>
     <a href="https://discord.com/invite/HhrNrHJPRb"><img src="https://img.shields.io/badge/discord-7289da.svg?style=flat-square&logo=discord" alt="discord" style="height: 20px;"></a>
     <a href="https://twitter.com/axolotl_ai"><img src="https://img.shields.io/twitter/follow/axolotl_ai?style=social" alt="twitter" style="height: 20px;"></a>
-    <a href="https://colab.research.google.com/github/axolotl-ai-cloud/axolotl/blob/main/examples/colab-notebooks/colab-axolotl-example.ipynb"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google-colab" style="height: 20px;"></a>
     <br/>
     <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/tests-nightly.yml/badge.svg" alt="tests-nightly">
     <img src="https://github.com/axolotl-ai-cloud/axolotl/actions/workflows/multi-gpu-e2e.yml/badge.svg" alt="multigpu-semi-weekly tests">
@@ -53,21 +49,20 @@
 
 ## ✨ Overview
 
-Axolotl is a free and open-source tool designed to streamline post-training and fine-tuning for the latest large language models (LLMs).
+Axolotl is a tool designed to streamline post-training for various AI models.
 
 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, and audio models like Voxtral with image, video, and audio support.
-- **Training Methods**: Full fine-tuning, LoRA, QLoRA, GPTQ, QAT, Preference Tuning (DPO, IPO, KTO, ORPO), RL (GRPO), 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.
+- **Multiple Model Support**: Train various models like LLaMA, Mistral, Mixtral, Pythia, and more. We are compatible with HuggingFace transformers causal language models.
+- **Training Methods**: Full fine-tuning, LoRA, QLoRA, GPTQ, QAT, Preference Tuning (DPO, IPO, KTO, ORPO), RL (GRPO), Multimodal, and Reward Modelling (RM) / Process Reward Modelling (PRM).
+- **Easy Configuration**: Re-use a single YAML file between dataset preprocess, training, evaluation, quantization, and inference.
 - **Performance Optimizations**: [Multipacking](https://docs.axolotl.ai/docs/multipack.html), [Flash Attention](https://github.com/Dao-AILab/flash-attention), [Xformers](https://github.com/facebookresearch/xformers), [Flex Attention](https://pytorch.org/blog/flexattention/), [Liger Kernel](https://github.com/linkedin/Liger-Kernel), [Cut Cross Entropy](https://github.com/apple/ml-cross-entropy/tree/main), [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
+## 🚀 Quick Start
 
 **Requirements**:
 
@@ -75,10 +70,6 @@ Features:
 - Python 3.11
 - PyTorch ≥2.6.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
 
 #### Using pip
@@ -164,7 +155,7 @@ If you use Axolotl in your research or projects, please cite it as follows:
 
 ```bibtex
 @software{axolotl,
-  title = {Axolotl: Open Source LLM Post-Training},
+  title = {Axolotl: Post-Training for AI Models},
   author = {{Axolotl maintainers and contributors}},
   url = {https://github.com/axolotl-ai-cloud/axolotl},
   license = {Apache-2.0},

_quarto.yml

@@ -153,7 +153,7 @@ quartodoc:
         - utils.distributed
         - utils.dict
         - utils.optimizers.adopt
-        - utils.data.streaming
+        - utils.data.pretraining
         - utils.data.sft
         - utils.quantization
     - title: Schemas
@@ -272,7 +272,6 @@ website:
           contents:
             - docs/batch_vs_grad.qmd
             - docs/dataset_preprocessing.qmd
-            - docs/streaming.qmd
             - docs/multipack.qmd
             - docs/mixed_precision.qmd
             - docs/optimizers.qmd

cicd/single_gpu.py

@@ -57,8 +57,7 @@ VOLUME_CONFIG = {
 }
 
 N_GPUS = int(os.environ.get("N_GPUS", 1))
-GPU_TYPE = os.environ.get("GPU_TYPE", "L40S")
-GPU_CONFIG = f"{GPU_TYPE}:{N_GPUS}"
+GPU_CONFIG = f"L40S:{N_GPUS}"
 
 
 def run_cmd(cmd: str, run_folder: str):

codecov.yml

@@ -12,7 +12,7 @@ coverage:
       default:
         # basic
         target: auto
-        threshold: 1%
+        threshold: 0%
         base: auto
         # advanced
         branches: null
@@ -27,7 +27,7 @@ coverage:
       default:
         # basic
         target: auto
-        threshold: 1%
+        threshold: 0%
         base: auto
         # advanced
         branches: null

docs/installation.qmd

@@ -134,7 +134,7 @@ For providers supporting Docker:
 
 ### Google Colab {#sec-colab}
 
-[![](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/axolotl-ai-cloud/axolotl/blob/main/examples/colab-notebooks/colab-axolotl-example.ipynb#scrollTo=msOCO4NRmRLa)
+Use our [example notebook](../examples/colab-notebooks/colab-axolotl-example.ipynb).
 
 ## Platform-Specific Instructions {#sec-platform-specific}
 

docs/multi-gpu.qmd

@@ -63,6 +63,15 @@ Start from Stage 1 -> Stage 2 -> Stage 3.
 
 :::
 
+::: {.callout-tip}
+
+Using ZeRO Stage 3 with Single-GPU training
+
+ZeRO Stage 3 can be used for training on a single GPU by manually setting the environment variables:
+`WORLD_SIZE=1 LOCAL_RANK=0 MASTER_ADDR=0.0.0.0 MASTER_PORT=29500`
+
+:::
+
 ## Fully Sharded Data Parallel (FSDP) {#sec-fsdp}
 
 ::: {.callout-note}

docs/multimodal.qmd

@@ -13,7 +13,6 @@ format:
 - [Pixtral](#sec-pixtral)
 - [Llava-1.5](#sec-llava-15)
 - [Mistral-Small-3.1](#sec-mistral-small-31)
-- [Magistral-Small-2509](#sec-magistral-small-2509)
 - [Voxtral](#sec-voxtral)
 - [Gemma-3](#sec-gemma-3)
 - [Gemma-3n](#sec-gemma-3n)
@@ -42,6 +41,7 @@ datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
+    field_messages: messages
 
 # (optional) if doing lora, only finetune the Language model,
 # leave the vision model and vision tower frozen
@@ -94,22 +94,10 @@ chat_template: llava
 
 ### Mistral-Small-3.1 {#sec-mistral-small-31}
 
-::: {.callout-tip}
-Please make sure to install vision lib via `pip install 'mistral-common[opencv]==1.8.5'`
-:::
-
 ```yaml
 base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
-```
 
-### Magistral-Small-2509 {#sec-magistral-small-2509}
-
-::: {.callout-tip}
-Please make sure to install vision lib via `pip install 'mistral-common[opencv]==1.8.5'`
-:::
-
-```yaml
-base_model: mistralai/Magistral-Small-2509
+chat_template: mistral_v7_tekken
 ```
 
 ### Voxtral {#sec-voxtral}

docs/qat.qmd

@@ -23,17 +23,10 @@ To enable QAT in axolotl, add the following to your configuration file:
 
 ```yaml
 qat:
-  activation_dtype: # Optional[str] = "int8". Fake quantization layout to use for activation quantization. Valid options are "int4", "int8", "float8"
-  weight_dtype: # Optional[str] = "int8". Fake quantization layout to use for weight quantization. Valid options are "int4", "fp8", and "nvfp4".
+  activation_dtype: # Optional[str] = "int8". Fake quantization layout to use for activation quantization. Valid options are "int4" and "int8"
+  weight_dtype: # Optional[str] = "int8". Fake quantization layout to use for weight quantization. Valid options are "int4" and "int8"
   group_size: # Optional[int] = 32. The number of elements in each group for per-group fake quantization
   fake_quant_after_n_steps: # Optional[int] = None. The number of steps to apply fake quantization after
 ```
 
-We support the following quantization schemas:
-- `Int4WeightOnly` (requires the `fbgemm-gpu` extra when installing Axolotl)
-- `Int8DynamicActivationInt4Weight`
-- `Float8DynamicActivationFloat8Weight`
-- `Float8DynamicActivationInt4Weight`
-- `NVFP4`
-
 Once you have finished training, you must quantize your model by using the same quantization configuration which you used to train the model with. You can use the [`quantize`](./quantize.qmd) command to do this.

docs/quantize.qmd

@@ -22,8 +22,8 @@ Quantization is configured using the `quantization` key in your configuration fi
 ```yaml
 base_model: # The path to the model to quantize.
 quantization:
-  activation_dtype: # Optional[str] = "int8". Fake quantization layout to use for activation quantization. Valid options are "int4", "int8", "float8"
-  weight_dtype: # Optional[str] = "int8". Fake quantization layout to use for weight quantization. Valid options are "int4", "fp8", and "nvfp4".
+  weight_dtype: # Optional[str] = "int8". Fake quantization layout to use for weight quantization. Valid options are uintX for X in [1, 2, 3, 4, 5, 6, 7], or int4, or int8
+  activation_dtype: # Optional[str] = "int8". Fake quantization layout to use for activation quantization. Valid options are "int4" and "int8"
   group_size: # Optional[int] = 32. The number of elements in each group for per-group fake quantization
   quantize_embedding: # Optional[bool] = False. Whether to quantize the embedding layer.
 
@@ -39,8 +39,9 @@ you used to train the model:
 # qat.yml
 qat:
   activation_dtype: int8
-  weight_dtype: int4
+  weight_dtype: int8
   group_size: 256
+  quantize_embedding: true
 
 output_dir: # The path to the output directory used during training where the final checkpoint has been saved.
 ```
@@ -50,11 +51,3 @@ axolotl quantize qat.yml
 ```
 
 This ensures that an identical quantization configuration is used to quantize the model as was used to train it.
-
-
-::: {.callout-note}
-
-If you have configured pushing to hub with `hub_model_id`, your model hub name will have the quantization schema appended to it,
-e.g. `axolotl-ai-cloud/qat-nvfp4-llama3B` will become `axolotl-ai-cloud/qat-nvfp4-llama3B-nvfp4w`
-
-:::

docs/reward_modelling.qmd

@@ -11,7 +11,6 @@ We support the reward modelling techniques supported by `trl`.
 ### (Outcome) Reward Models
 
 Outcome reward models are trained using data which contains preference annotations for an entire interaction between the user and model (e.g. rather than per-turn or per-step).
-For improved training stability, you can use the `center_rewards_coefficient` parameter to encourage mean-zero reward outputs ([see TRL docs](https://huggingface.co/docs/trl/v0.10.1/en/reward_trainer#centering-rewards)).
 
 ```yaml
 base_model: google/gemma-2-2b

docs/streaming.qmd

@@ -1,120 +0,0 @@
----
-title: Streaming Datasets
-description: How to use streaming mode for large-scale datasets and memory-efficient training
-order: 10
----
-
-Streaming enables memory-efficient training with large datasets by loading data
-incrementally rather than loading the entire dataset into memory at once.
-
-Use streaming when:
-
-- Your dataset is too large to fit in memory (e.g. when you're doing pretraining with massive text corpora)
-- You want to start training immediately without preprocessing the entire dataset
-
-Streaming works with both remote and locally stored datasets!
-
-::: {.callout-note}
-Streaming currently only supports a single dataset. Multi-dataset support will be added soon.
-:::
-
-
-## Configuration
-
-### Basic Streaming
-
-Enable streaming mode by setting the `streaming` flag:
-
-```yaml
-streaming: true
-```
-
-### Pretraining with Streaming
-
-For pretraining tasks, streaming is automatically enabled when using `pretraining_dataset`:
-
-```yaml
-pretraining_dataset:
-  - path: HuggingFaceFW/fineweb-edu
-    type: pretrain
-    text_column: text
-    split: train
-
-# Optionally, enable sample packing
-streaming_multipack_buffer_size: 10000
-sample_packing: true
-```
-
-### SFT with Streaming
-
-For supervised fine-tuning with streaming:
-
-```yaml
-streaming: true
-datasets:
-  - path: tatsu-lab/alpaca
-    type: alpaca
-    split: train
-
-# Optionally, enable sample packing
-streaming_multipack_buffer_size: 10000
-sample_packing: true
-```
-
-## Configuration Options
-
-### `streaming_multipack_buffer_size`
-
-Controls the buffer size for multipack streaming (default: 10,000). This determines how
-many samples are buffered before packing. Larger buffers can improve packing efficiency
-but use more memory.
-
-### `shuffle_merged_datasets`
-
-When enabled, shuffles the streaming dataset using the buffer. This requires additional
-memory for the shuffle buffer.
-
-## Sample Packing with Streaming
-
-Sample packing is supported for streaming datasets. When enabled, multiple samples are
-packed into a single sequence to maximize GPU utilization:
-
-```yaml
-sample_packing: true
-streaming_multipack_buffer_size: 10000
-
-# For SFT: attention is automatically isolated between packed samples
-# For pretraining: control with pretrain_multipack_attn
-pretrain_multipack_attn: true  # prevent cross-attention between packed samples
-```
-
-For more information, see our [documentation](multipack.qmd) on multipacking.
-
-## Important Considerations
-
-### Memory Usage
-
-While streaming reduces memory usage compared to loading entire datasets, you still need
-to consider:
-
-- You can control the memory usage by adjusting `streaming_multipack_buffer_size`
-- Sample packing requires buffering multiple samples
-- Shuffling requires additional memory for the shuffle buffer
-
-### Performance
-
-- Streaming may have slightly higher latency compared to preprocessed datasets, as samples are processed on-the-fly
-- Network speed and disk read speed are important when streaming from remote sources or a local dataset, respectively
-- Consider using `axolotl preprocess` for smaller or more frequently used datasets
-
-### Evaluation Datasets
-
-Evaluation datasets are not streamed to ensure consistent evaluation metrics. They're
-loaded normally even when training uses streaming.
-
-## Examples
-
-See the `examples/streaming/` directory for complete configuration examples:
-
-- `pretrain.yaml`: Pretraining with streaming dataset
-- `sft.yaml`: Supervised fine-tuning with streaming

examples/apertus/README.md

@@ -1,110 +0,0 @@
-# Finetune Swiss-AI's Apertus with Axolotl
-
-[Apertus](https://huggingface.co/collections/swiss-ai/apertus-llm-68b699e65415c231ace3b059) is a family of opensource models trained by Swiss-ai.
-
-This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as Apertus is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
-
-    Here is an example of how to install from main for pip:
-
-```bash
-# Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-git clone https://github.com/axolotl-ai-cloud/axolotl.git
-cd axolotl
-
-pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
-pip3 install --no-build-isolation -e '.[flash-attn]'
-
-# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
-python scripts/cutcrossentropy_install.py | sh
-```
-
-2. (Optional, highly recommended) Install XIELU CUDA
-
-```bash
-## Recommended for reduced VRAM and faster speeds
-
-# Point to CUDA toolkit directory
-# For those using our Docker image, use the below path.
-export CUDA_HOME=/usr/local/cuda
-
-pip3 install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
-```
-
-For any installation errors, see [XIELU Installation Issues](#xielu-installation-issues)
-
-3. Run the finetuning example:
-
-```bash
-axolotl train examples/apertus/apertus-8b-qlora.yaml
-```
-
-This config uses about 8.7 GiB VRAM.
-
-Let us know how it goes. Happy finetuning! 🚀
-
-### Tips
-
-- For inference, the official Apertus team recommends `top_p=0.9` and `temperature=0.8`.
-- You can instead use full paremter fine-tuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
-- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
-
-### XIELU Installation Issues
-
-#### `ModuleNotFoundError: No module named 'torch'`
-
-Please check these one by one:
-- Running in correct environment
-- Env has PyTorch installed
-- CUDA toolkit is at `CUDA_HOME`
-
-If those didn't help, please try the below solutions:
-
-1. Pass env for CMAKE and try install again:
-
-    ```bash
-    Python_EXECUTABLE=$(which python) pip3 install git+https://github.com/nickjbrowning/XIELU@59d6031 --no-build-isolation --no-deps
-    ```
-
-2. Git clone the repo and manually hardcode python path:
-
-    ```bash
-    git clone https://github.com/nickjbrowning/XIELU
-    cd xielu
-    git checkout 59d6031
-
-    cd xielu
-    nano CMakeLists.txt  # or vi depending on your preference
-    ```
-
-    ```diff
-    execute_process(
-    -    COMMAND ${Python_EXECUTABLE} -c "import torch.utils; print(torch.utils.cmake_prefix_path)"
-    +    COMMAND /root/miniconda3/envs/py3.11/bin/python -c "import torch.utils; print(torch.utils.cmake_prefix_path)"
-        RESULT_VARIABLE TORCH_CMAKE_PATH_RESULT
-        OUTPUT_VARIABLE TORCH_CMAKE_PATH_OUTPUT
-        ERROR_VARIABLE TORCH_CMAKE_PATH_ERROR
-    )
-    ```
-
-    ```bash
-    pip3 install . --no-build-isolation --no-deps
-    ```
-
-## Optimization Guides
-
-- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
-- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
-- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
-
-## Related Resources
-
-- [Apertus Tech Report](https://github.com/swiss-ai/apertus-tech-report/blob/main/Apertus_Tech_Report.pdf)
-- [Axolotl Docs](https://docs.axolotl.ai)
-- [Axolotl Website](https://axolotl.ai)
-- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
-- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/apertus/apertus-8b-qlora.yaml

@@ -1,64 +0,0 @@
-base_model: swiss-ai/Apertus-8B-Instruct-2509
-
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: false
-load_in_4bit: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/lora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/arcee/README.md

@@ -19,9 +19,6 @@ cd axolotl
 
 pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation -e '.[flash-attn]'
-
-# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
-python scripts/cutcrossentropy_install.py | sh
 ```
 
 2. Run the finetuning example:

examples/archived/deepcoder/deepcoder-14B-preview-lora.yml

@@ -9,6 +9,10 @@ strict: false
 datasets:
   - path: fozziethebeat/alpaca_messages_2k_test
     type: chat_template
+    field_messages: messages
+    message_property_mappings:
+      role: role
+      content: content
 
 dataset_prepared_path:
 val_set_size: 0.05

examples/colab-notebooks/colab-axolotl-example.ipynb

@@ -40,7 +40,7 @@
     "%%capture\n",
     "# This step can take ~5-10 minutes to install dependencies\n",
     "!pip install --no-build-isolation axolotl[flash-attn]>=0.9.1\n",
-    "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c5aa3ef\""
+    "!pip install \"cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5\""
    ]
   },
   {
@@ -176,8 +176,8 @@
     }
    ],
    "source": [
-    "from axolotl.cli.config import load_cfg\n",
     "from axolotl.utils.dict import DictDefault\n",
+    "from axolotl.cli.config import load_cfg\n",
     "\n",
     "# Axolotl provides full control and transparency over model and training configuration\n",
     "config = DictDefault(\n",
@@ -251,10 +251,10 @@
    },
    "outputs": [],
    "source": [
-    "from axolotl.utils import set_pytorch_cuda_alloc_conf\n",
+    "from axolotl.utils import patch_optimized_env\n",
     "\n",
-    "# Set \"PYTORCH_CUDA_ALLOC_CONF\" env to save memory\n",
-    "set_pytorch_cuda_alloc_conf()"
+    "# speedup downloads from HF 🤗 and set \"PYTORCH_CUDA_ALLOC_CONF\" env to save memory\n",
+    "patch_optimized_env()"
    ]
   },
   {

examples/deepcogito/cogito-v1-preview-llama-3B-lora.yml

@@ -9,6 +9,10 @@ strict: false
 datasets:
   - path: fozziethebeat/alpaca_messages_2k_test
     type: chat_template
+    field_messages: messages
+    message_property_mappings:
+      role: role
+      content: content
 
 dataset_prepared_path:
 val_set_size: 0.05

examples/deepcogito/cogito-v1-preview-qwen-14B-lora.yml

@@ -9,6 +9,10 @@ strict: false
 datasets:
   - path: fozziethebeat/alpaca_messages_2k_test
     type: chat_template
+    field_messages: messages
+    message_property_mappings:
+      role: role
+      content: content
 
 dataset_prepared_path:
 val_set_size: 0.05

examples/devstral/README.md

@@ -20,13 +20,7 @@ pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage
-
-```bash
-python scripts/cutcrossentropy_install.py | sh
-```
-
-3. Run the finetuning example:
+2. Run the finetuning example:
 
 ```bash
 axolotl train examples/devstral/devstral-small-qlora.yml

examples/gemma3/gemma-3-270m-qlora.yml

@@ -1,68 +0,0 @@
-base_model: google/gemma-3-270m-it
-# optionally might have model_type or tokenizer_type
-model_type: AutoModelForCausalLM
-tokenizer_type: AutoTokenizer
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-# gemma3 doesn't seem to play nice with ddp
-ddp_find_unused_parameters: true
-
-load_in_8bit: false
-load_in_4bit: true
-
-# huggingface repo
-chat_template: gemma3
-eot_tokens:
-  - <end_of_turn>
-datasets:
-  - path: cgato/SlimOrcaDedupCleaned
-    type: chat_template
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-adapter: qlora
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: false
-
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch:
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:

examples/gemma3/gemma-3-4b-vision-qlora.yml

@@ -18,7 +18,7 @@ datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
-
+    field_messages: messages
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.01
 output_dir: ./outputs/out

examples/gemma3n/README.md

@@ -23,15 +23,7 @@ pip3 install timm==1.0.17
 pip3 install librosa==0.11.0
 ```
 
-3. Download sample dataset files
-
-```bash
-# for text + vision + audio only
-wget https://huggingface.co/datasets/Nanobit/text-vision-audio-2k-test/resolve/main/African_elephant.jpg
-wget https://huggingface.co/datasets/Nanobit/text-vision-audio-2k-test/resolve/main/En-us-African_elephant.oga
-```
-
-4. Run the finetuning example:
+3. Run the finetuning example:
 
 ```bash
 # text only

examples/gpt-oss/README.md

@@ -106,16 +106,6 @@ See [Nanobit/text-tools-2k-test](https://huggingface.co/datasets/Nanobit/text-to
 
 Refer to [our docs](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#using-tool-use) for more info.
 
-### Thinking and chat_template masking conflict
-
-OpenAI’s Harmony template hides `thinking` in all non-final turns, which conflicts with Axolotl’s `chat_template` masking.
-
-If your dataset has `thinking` content mid-turn, there are two paths we recommend:
-
-- Train only on the last turn. This can be accomplished via chat_template's [train on last doc](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#training-on-last-message).
-
-- Adjust your dataset to only have `thinking` content in the last turn.
-
 ### TIPS
 
 - Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).

examples/hunyuan/README.md

@@ -1,85 +0,0 @@
-# Finetune HunYuan with Axolotl
-
-Tencent released a family of opensource models called HunYuan with varying parameter scales of 0.5B, 1.8B, 4B, and 7B scale for both Pre-trained and Instruct variants. The models can be found at [HuggingFace](https://huggingface.co/collections/tencent/hunyuan-dense-model-6890632cda26b19119c9c5e7). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as HunYuan is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
-
-    Here is an example of how to install from main for pip:
-
-```bash
-# Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-git clone https://github.com/axolotl-ai-cloud/axolotl.git
-cd axolotl
-
-pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
-pip3 install --no-build-isolation -e '.[flash-attn]'
-
-# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
-python scripts/cutcrossentropy_install.py | sh
-```
-
-2. Run the finetuning example:
-
-```bash
-axolotl train examples/hunyuan/hunyuan-v1-dense-qlora.yaml
-```
-
-This config uses about 4.7 GB VRAM.
-
-Let us know how it goes. Happy finetuning! 🚀
-
-### Dataset
-
-HunYuan Instruct models can choose to enter a slow think or fast think pattern. For best performance on fine-tuning their Instruct models, your dataset should be adjusted to match their pattern.
-
-```python
-# fast think pattern
-messages = [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "/no_think What color is the sun?" },
-    {"role": "assistant", "content": "<think>\n\n</think>\n<answer>\nThe sun is yellow.\n</answer>"}
-]
-
-# slow think pattern
-messages = [
-    {"role": "system", "content": "You are a helpful assistant."},
-    {"role": "user", "content": "/no_think What color is the sun?" },
-    {"role": "assistant", "content": "<think>\nThe user is asking about the color of the sun. I need to ...\n</think>\n<answer>\nThe sun is yellow.\n</answer>"}
-]
-```
-
-### TIPS
-
-- For inference, the official Tencent team recommends
-
-```json
-
-{
-  "do_sample": true,
-  "top_k": 20,
-  "top_p": 0.8,
-  "repetition_penalty": 1.05,
-  "temperature": 0.7
-}
-
-```
-
-- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
-- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
-
-## Optimization Guides
-
-- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
-- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
-- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
-
-## Related Resources
-
-- [Tencent HunYuan Blog](https://hunyuan.tencent.com/)
-- [Axolotl Docs](https://docs.axolotl.ai)
-- [Axolotl Website](https://axolotl.ai)
-- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
-- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/hunyuan/hunyuan-v1-dense-qlora.yaml

@@ -1,64 +0,0 @@
-base_model: tencent/Hunyuan-0.5B-Instruct
-
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: false
-load_in_4bit: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/lora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-lora_target_modules:
-  - gate_proj
-  - down_proj
-  - up_proj
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/llama-3/3b-qat-fsdp2.yaml

@@ -15,18 +15,20 @@ liger_glu_activation: true
 liger_layer_norm: true
 liger_fused_linear_cross_entropy: true
 
-
 datasets:
   - path: yahma/alpaca-cleaned
     type: alpaca
-    split: train[:95%]
 
 output_dir: ./outputs/qat_out/
-dataset_prepared_path: ./outputs/qat_out/dataset_prepared
 
-sample_packing: false
-sequence_len: 8192
-flash_attention: true
+sample_packing: true
+
+sequence_len: 512
+
+flex_attention: true
+flex_attn_compile_kwargs:
+  dynamic: false
+  mode: max-autotune-no-cudagraphs
 
 qat:
   activation_dtype: int8
@@ -65,7 +67,7 @@ fsdp:
 fsdp_config:
   fsdp_version: 2
   fsdp_offload_params: false
-  fsdp_cpu_ram_efficient_loading: false
+  fsdp_cpu_ram_efficient_loading: true
   fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
   fsdp_transformer_layer_cls_to_wrap: LlamaDecoderLayer
   fsdp_state_dict_type: FULL_STATE_DICT
@@ -74,6 +76,6 @@ fsdp_config:
   fsdp_activation_checkpointing: true
 
 special_tokens:
-  pad_token: <|finetune_right_pad_id|>
+  pad_token: <|end_of_text|>
 
 # save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/llama-3/3b-qat-nvfp4.yaml

@@ -1,64 +0,0 @@
-base_model: meta-llama/Llama-3.2-3B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-load_in_8bit: false
-load_in_4bit: false
-strict: false
-
-plugins:
-  - axolotl.integrations.liger.LigerPlugin
-
-liger_rope: true
-liger_rms_norm: true
-liger_glu_activation: true
-liger_layer_norm: true
-liger_fused_linear_cross_entropy: true
-
-datasets:
-  - path: yahma/alpaca-cleaned
-    type: alpaca
-    split: train[:95%]
-
-output_dir: ./outputs/qat_out/
-dataset_prepared_path: ./outputs/dataset_prepared
-
-sequence_len: 8192
-flash_attention: true
-
-qat:
-  activation_dtype: nvfp4
-  weight_dtype: nvfp4
-  group_size: 16 # only group_size of 16 is supported with nvfp4
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_checkpointing: true
-gradient_accumulation_steps: 1
-micro_batch_size: 64
-num_epochs: 1
-optimizer: adamw_torch_fused
-
-cosine_constant_lr_ratio: 0
-cosine_min_lr_ratio: 1.0
-learning_rate: 2e-5
-save_only_model: true
-bf16: true
-
-resume_from_checkpoint:
-logging_steps: 1
-
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-warmup_ratio: 0.1
-weight_decay: 0.0
-
-special_tokens:
-  pad_token: <|finetune_right_pad_id|>
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/llama-3/diffusion/pretrain-1b.yaml

@@ -1,56 +0,0 @@
-base_model: meta-llama/Llama-3.2-1B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-pretraining_dataset:
-  - path: wikitext
-    name: wikitext-103-raw-v1
-    type: completion
-    field: text
-
-plugins:
-  - axolotl.integrations.diffusion.DiffusionPlugin
-
-diffusion:
-  noise_schedule: cosine
-  min_mask_ratio: 0.15
-  max_mask_ratio: 0.85
-  num_diffusion_steps: 128
-  eps: 5e-4
-  importance_weighting: true
-  mask_token_id: 128002
-  generate_samples: true
-  generation_interval: 250
-
-output_dir: ./outputs/model-out
-
-sequence_len: 512
-sample_packing: true
-
-gradient_accumulation_steps: 8
-micro_batch_size: 4
-max_steps: 10000
-warmup_ratio: 0.1
-
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 3e-4
-sdp_attention: true
-
-bf16: auto
-tf32: true
-
-logging_steps: 1
-save_strategy: steps
-save_steps: 1000
-
-special_tokens:
-  pad_token: "<|end_of_text|>"
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/llama-3/diffusion/sft-1b.yaml

@@ -1,59 +0,0 @@
-base_model: meta-llama/Llama-3.2-1B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-datasets:
-  - path: teknium/GPT4-LLM-Cleaned
-    type: alpaca
-val_set_size: 0.05
-
-plugins:
-  - axolotl.integrations.diffusion.DiffusionPlugin
-
-diffusion:
-  noise_schedule: cosine
-  min_mask_ratio: 0.1
-  max_mask_ratio: 0.9
-  num_diffusion_steps: 128
-  eps: 1e-3
-  importance_weighting: true
-  mask_token_id: 128002
-  generate_samples: true
-  generation_interval: 250
-
-output_dir: ./outputs/model-out
-
-sequence_len: 512
-sample_packing: true
-eval_sample_packing: true
-
-gradient_accumulation_steps: 4
-micro_batch_size: 4
-num_epochs: 1
-warmup_steps: 0.1
-
-optimizer: adamw_8bit
-lr_scheduler: cosine
-learning_rate: 1e-5
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-sdp_attention: true
-
-logging_steps: 1
-save_strategy: best
-eval_strategy: epoch
-
-special_tokens:
-  pad_token: "<|end_of_text|>"
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/llama-3/instruct-lora-8b.yml

@@ -12,6 +12,15 @@ chat_template: llama3
 datasets:
   - path: fozziethebeat/alpaca_messages_2k_test
     type: chat_template
+    field_messages: messages
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      user:
+        - user
+      assistant:
+        - assistant
 
 dataset_prepared_path:
 val_set_size: 0.05

examples/llama-4/do-no-use-fa2/scout-vision-qlora-fsdp.yaml

@@ -46,6 +46,7 @@ datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
+    field_messages: messages
 
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0

examples/llama-4/scout-vision-qlora-fsdp2-flex.yaml

@@ -45,6 +45,7 @@ datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
+    field_messages: messages
 
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0

examples/magistral/README.md

@@ -1,10 +1,10 @@
 # Finetune Magistral Small with Axolotl
 
-Magistral Small is a 24B parameter opensource model from MistralAI found on HuggingFace at [2506](https://huggingface.co/mistralai/Magistral-Small-2506), [2507](https://huggingface.co/mistralai/Magistral-Small-2507) (see [Thinking](#thinking)), and [2509](https://huggingface.co/mistralai/Magistral-Small-2509) (see [Vision](#vision)). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
+Magistral Small is a 24B parameter opensource model from MistralAI found on HuggingFace at [2506](https://huggingface.co/mistralai/Magistral-Small-2506) and [2507](https://huggingface.co/mistralai/Magistral-Small-2507) (see [Thinking](#thinking)). This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
 
 MistralAI has also released a proprietary medium-sized version called Magistral Medium.
 
-Thanks to the team at MistralAI for giving us early access to prepare for these releases.
+Thanks to the team at MistralAI for giving us early access to prepare for this release.
 
 ## Getting started
 
@@ -18,13 +18,7 @@ pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
 pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 ```
 
-2. Install [Cut Cross Entropy](https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy) to reduce training VRAM usage
-
-```bash
-python scripts/cutcrossentropy_install.py | sh
-```
-
-3. Run the finetuning example:
+2. Run the finetuning example:
 
 ```bash
 axolotl train examples/magistral/magistral-small-qlora.yaml
@@ -36,17 +30,29 @@ Let us know how it goes. Happy finetuning! 🚀
 
 ### Thinking
 
-MistralAI has released their [2507](https://huggingface.co/mistralai/Magistral-Small-2507) model with thinking capabilities, enabling Chain-of-Thought reasoning with explicit thinking steps.
+MistralAI has released their [2507](https://huggingface.co/mistralai/Magistral-Small-2507) model with thinking capabilities. The model requires the multi-content dataset format with support for an extra `role: thinking` within system and assistant messages.
 
-📚 **[See the Thinking fine-tuning guide →](./think/README.md)**
+Example format:
 
-### Vision
+```json
+{
+    "messages": [
+        {"role": "system", "content": [{ "type": "text", "text": "{SYSTEM_PROMPT}"}]},
+        {"role": "user", "content": [{ "type": "text", "text": "..."}]},
+        {"role": "assistant", "content": [{ "type": "thinking", "thinking": "..."}, { "type": "text", "text": "..." }]},
+    ],
+}
+```
 
-MistralAI has released their [2509](https://huggingface.co/mistralai/Magistral-Small-2509) model with vision capabilities.
+Example config: `./magistral-small-think-qlora.yaml`.
 
-📚 **[See the Vision fine-tuning guide →](./vision/README.md)**
+The `thinking` section also supports an optional arg `closed: bool` (`True` default) which controls adding the closing `[/THINK]` tag.
 
-### Tips
+Limitations:
+- You cannot mix `content: str` with `content: list[dict]` as the `dataset.load_dataset` may complain about different types for `content` key.
+- This mode does not work with custom `train_detail` and `training` at the moment.
+
+### TIPS
 
 - We recommend adding the same/similar SystemPrompt that the model is tuned for. You can find this within the repo's files titled `SYSTEM_PROMPT.txt`.
 - For inference, the official MistralAI team recommends `top_p: 0.95` and `temperature: 0.7` with `max_tokens: 40960`.
@@ -77,5 +83,5 @@ In addition, we do not support overriding tokens yet.
 
 ## Future Work
 
-- Add parity to Preference Tuning, RL, etc.
+- Add parity to Preference Tuning, RL, Multi-modal, etc.
 - Add parity to other tokenizer configs like overriding tokens.

examples/magistral/think/README.md

@@ -1,73 +0,0 @@
-# Magistral Small Thinking Fine-tuning
-
-This guide covers fine-tuning [Magistral Small 2507](https://huggingface.co/mistralai/Magistral-Small-2507) with thinking capabilities using Axolotl. The thinking model enables explicit Chain-of-Thought reasoning with separate thinking and response sections.
-
-## Prerequisites
-
-Before starting, ensure you have:
-- Installed Axolotl (see [main README](../README.md))
-
-## Getting Started
-
-Run the thinking model fine-tuning:
-
-```bash
-axolotl train magistral-small-think-qlora.yaml
-```
-
-This config uses about 19.1 GiB VRAM.
-
-### Tips
-
-- Dataset uses multi-content format with `type: thinking` support. See [Dataset Format](#dataset-format) below.
-- You cannot mix `content: str` and `content: list[dict]`, otherwise, dataset loading will fail. Keep it consistent.
-
-## Dataset Format
-
-The thinking model requires the multi-content dataset format with support for an extra `role: thinking` within system and assistant messages.
-
-Example format:
-
-```json
-{
-    "messages": [
-        {
-            "role": "system",
-            "content": [
-                { "type": "text", "text": "{SYSTEM_PROMPT}"}
-            ]
-        },
-        {
-            "role": "user",
-            "content": [
-                { "type": "text", "text": "Solve this step by step: What is 15% of 240?"}
-            ]
-        },
-        {
-            "role": "assistant",
-            "content": [
-                {
-                    "type": "thinking",
-                    "thinking": "I need to calculate 15% of 240. First, I'll convert 15% to decimal: 0.15. Then multiply: 0.15 × 240 = 36."
-                },
-                {
-                    "type": "text",
-                    "text": "To find 15% of 240, I'll multiply 240 by 0.15:\n\n240 × 0.15 = 36\n\nTherefore, 15% of 240 is 36."
-                }
-            ]
-        }
-    ]
-}
-```
-
-### Advanced Options
-
-The `thinking` section supports an optional `closed` parameter:
-
-```json
-{
-    "type": "thinking",
-    "thinking": "Internal reasoning here...",
-    "closed": true  // Default: true, controls adding the closing [/THINK] tag
-}
-```

examples/magistral/vision/README.md

@@ -1,60 +0,0 @@
-# Magistral Small Vision Fine-tuning
-
-This guide covers fine-tuning [Magistral Small 2509](https://huggingface.co/mistralai/Magistral-Small-2509) with vision capabilities using Axolotl.
-
-## Prerequisites
-
-Before starting, ensure you have:
-- Installed Axolotl from source (see [main README](../README.md#getting-started))
-
-## Getting started
-
-1. Install the required vision lib:
-    ```bash
-    pip install 'mistral-common[opencv]==1.8.5'
-    ```
-
-2. Download the example dataset image:
-   ```bash
-   wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
-   ```
-
-3. Run the fine-tuning:
-   ```bash
-   axolotl train magistral-small-vision-24B-qlora.yml
-   ```
-
-This config uses about 17GiB VRAM.
-
-WARNING: The loss and grad norm will be much higher than normal at first. We suspect this to be inherent to the model as of the moment. If anyone would like to submit a fix for this, we are happy to take a look.
-
-### Tips
-
-Key differences from text-only model:
-- `max_tokens: 131072` for inference
-- Multi-modal dataset format required
-- Sample packing not supported
-
-## Dataset Format
-
-The vision model requires multi-modal dataset format as documented [here](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
-
-One exception is that, passing `"image": PIL.Image` is not supported. MistralTokenizer only supports `path`, `url`, and `base64` for now.
-
-Example:
-```json
-{
-    "messages": [
-        {"role": "system", "content": [{ "type": "text", "text": "{SYSTEM_PROMPT}"}]},
-        {"role": "user", "content": [
-            { "type": "text", "text": "What's in this image?"},
-            {"type": "image", "path": "path/to/image.jpg" }
-        ]},
-        {"role": "assistant", "content": [{ "type": "text", "text": "..." }]},
-    ],
-}
-```
-
-## Limitations
-
-- Sample Packing is not supported for multi-modality training currently.

examples/magistral/vision/magistral-small-vision-24B-qlora.yml

@@ -1,64 +0,0 @@
-base_model: mistralai/Magistral-Small-2509
-processor_type: AutoProcessor
-
-# Enable to use mistral-common tokenizer
-tokenizer_use_mistral_common: true
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_4bit: true
-
-# these 3 lines are needed for now to handle vision chat templates w images
-skip_prepare_dataset: true
-remove_unused_columns: false
-sample_packing: false
-
-# sample dataset below requires downloading image in advance
-# wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
-datasets:
-  - path: Nanobit/text-vision-2k-test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.01
-output_dir: ./outputs/out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 1
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: true
-fp16:
-tf32: true
-
-gradient_checkpointing: true
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/mistral/mistral-small-3.1-24B-lora.yml

@@ -1,9 +1,6 @@
 base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
 processor_type: AutoProcessor
 
-# Enable to use mistral-common tokenizer
-tokenizer_use_mistral_common: true
-
 load_in_8bit: true
 
 # these 3 lines are needed for now to handle vision chat templates w images
@@ -11,12 +8,12 @@ skip_prepare_dataset: true
 remove_unused_columns: false
 sample_packing: false
 
-# sample dataset below requires downloading image in advance
-# wget https://huggingface.co/datasets/Nanobit/text-vision-2k-test/resolve/main/African_elephant.jpg
+chat_template: mistral_v7_tekken
 datasets:
-  - path: Nanobit/text-vision-2k-test
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
-
+    split: train[:1%]
+    field_messages: messages
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.01
 output_dir: ./outputs/out
@@ -51,7 +48,8 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: true
+# flash_attention: false # PixtralVisionModel does not support Flash Attention 2.0 yet.
+sdp_attention: true
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/phi/lora-3.5.yaml

@@ -12,6 +12,15 @@ chat_template: phi_3
 datasets:
   - path: fozziethebeat/alpaca_messages_2k_test
     type: chat_template
+    field_messages: messages
+    message_property_mappings:
+      role: role
+      content: content
+    roles:
+      user:
+        - user
+      assistant:
+        - assistant
 
 dataset_prepared_path:
 val_set_size: 0.05

examples/pixtral/lora-12b.yml

@@ -45,7 +45,8 @@ tf32: true
 
 gradient_checkpointing: true
 logging_steps: 1
-flash_attention: true
+# flash_attention:  # PixtralVisionModel does not support Flash Attention 2.0 yet
+sdp_attention: true
 
 warmup_ratio: 0.1
 evals_per_epoch: 1

examples/qwen2-vl/lora-7b.yaml

@@ -11,7 +11,7 @@ datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
-
+    field_messages: messages
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
 output_dir: ./outputs/out

examples/qwen2_5-vl/lora-7b.yaml

@@ -11,7 +11,7 @@ datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
-
+    field_messages: messages
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
 output_dir: ./outputs/out

examples/qwen3-next/README.md

@@ -1,64 +0,0 @@
-# Finetune Qwen3-Next with Axolotl
-
-[Qwen3-Next](https://huggingface.co/collections/Qwen/qwen3-next-68c25fd6838e585db8eeea9d) represents the next-generation foundation models optimized for extreme context length and large-scale parameter efficiency. The series introduces architectural innovations including Hybrid Attention (Gated DeltaNet + Gated Attention), High-Sparsity MoE with 1:50 activation ratio, and Multi-Token Prediction for enhanced performance and inference acceleration.
-
-This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as Qwen3-Next is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
-
-    Here is an example of how to install from main for pip:
-
-```bash
-# Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-git clone https://github.com/axolotl-ai-cloud/axolotl.git
-cd axolotl
-
-pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
-pip3 install --no-build-isolation -e '.[flash-attn]'
-
-# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
-python scripts/cutcrossentropy_install.py | sh
-```
-
-2. Install Qwen3-Next transformers commit
-```bash
-pip3 uninstall -y transformers && pip3 install "git+https://github.com/huggingface/transformers.git@b9282355bea846b54ed850a066901496b19da654"
-```
-
-3. Install FLA for improved performance
-```bash
-pip3 uninstall -y causal-conv1d && pip3 install flash-linear-attention==0.3.2
-```
-
-4. Run the finetuning example:
-
-```bash
-axolotl train examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml
-```
-
-This config uses about 41.7 GiB VRAM.
-
-Let us know how it goes. Happy finetuning! 🚀
-
-### TIPS
-
-- For inference, you can experiment with `temperature: 0.7`, `top_p: 0.8`, `top_k: 20`, and `min_p: 0`.
-- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config. See [Multi-GPU](#optimization-guides) section below.
-- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
-
-## Optimization Guides
-
-- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
-- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
-- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
-
-## Related Resources
-
-- [Qwen3-Next Blog](https://qwenlm.github.io/blog/qwen3_next/)
-- [Axolotl Docs](https://docs.axolotl.ai)
-- [Axolotl Website](https://axolotl.ai)
-- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
-- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/qwen3-next/qwen3-next-80b-a3b-qlora.yaml

@@ -1,60 +0,0 @@
-base_model: Qwen/Qwen3-Next-80B-A3B-Instruct
-
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: false
-load_in_4bit: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/lora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 16
-lora_alpha: 8
-lora_dropout: 0.05
-lora_target_modules:
-  - q_proj
-  - v_proj
-  - k_proj
-  - o_proj
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/qwen3/reward-model.yaml

@@ -1,44 +0,0 @@
-base_model: Skywork/Skywork-Reward-V2-Qwen3-8B
-model_type: AutoModelForSequenceClassification
-num_labels: 1
-
-reward_model: true
-center_rewards_coefficient: 0.01  # Incentivize mean-zero rewards for improved stability
-chat_template: qwen3
-datasets:
-  - path: argilla/distilabel-intel-orca-dpo-pairs
-    type: bradley_terry.chat_template
-
-val_set_size: 0.0
-output_dir: ./outputs/out
-
-sequence_len: 8192
-sample_packing: false
-eval_sample_packing: false
-pad_to_sequence_len: true
-
-deepspeed: deepspeed_configs/zero1.json
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-eval_batch_size: 1
-num_epochs: 3
-optimizer: adamw_bnb_8bit
-lr_scheduler: linear
-learning_rate: 0.00002
-
-bf16: true
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-warmup_ratio: 0.1
-logging_steps: 1
-weight_decay: 0.01

examples/seed-oss/README.md

@@ -1,54 +0,0 @@
-# Finetune ByteDance's Seed-OSS with Axolotl
-
-[Seed-OSS](https://huggingface.co/collections/ByteDance-Seed/seed-oss-68a609f4201e788db05b5dcd) are a series of 36B parameter open source models trained by ByteDance's Seed Team.
-
-This guide shows how to fine-tune it with Axolotl with multi-turn conversations and proper masking.
-
-## Getting started
-
-1. Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html). You need to install from main as Seed-OSS is only on nightly or use our latest [Docker images](https://docs.axolotl.ai/docs/docker.html).
-
-    Here is an example of how to install from main for pip:
-
-```bash
-# Ensure you have Pytorch installed (Pytorch 2.6.0 min)
-git clone https://github.com/axolotl-ai-cloud/axolotl.git
-cd axolotl
-
-pip3 install packaging==23.2 setuptools==75.8.0 wheel ninja
-pip3 install --no-build-isolation -e '.[flash-attn]'
-
-# Install Cut Cross Entropy
-python scripts/cutcrossentropy_install.py | sh
-```
-
-2. Run the finetuning example:
-
-```bash
-axolotl train examples/seed-oss/seed-oss-36b-qlora.yaml
-```
-
-This config uses about 27.7 GiB VRAM.
-
-Let us know how it goes. Happy finetuning! 🚀
-
-### TIPS
-
-- For inference, the official Seed Team recommends `top_p=0.95` and `temperature=1.1`.
-- You can run a full finetuning by removing the `adapter: qlora` and `load_in_4bit: true` from the config.
-- Read more on how to load your own dataset at [docs](https://docs.axolotl.ai/docs/dataset_loading.html).
-- The dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
-
-## Optimization Guides
-
-- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
-- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
-- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
-
-## Related Resources
-
-- [ByteDance Seed Website](https://seed.bytedance.com/)
-- [Axolotl Docs](https://docs.axolotl.ai)
-- [Axolotl Website](https://axolotl.ai)
-- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
-- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/seed-oss/seed-oss-36b-qlora.yaml

@@ -1,56 +0,0 @@
-base_model: ByteDance-Seed/Seed-OSS-36B-Instruct
-
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-
-load_in_8bit: false
-load_in_4bit: true
-
-datasets:
-  - path: fozziethebeat/alpaca_messages_2k_test
-    type: chat_template
-
-dataset_prepared_path: last_run_prepared
-val_set_size: 0.1
-output_dir: ./outputs/lora-out
-
-adapter: qlora
-lora_model_dir:
-
-sequence_len: 2048
-sample_packing: true
-
-lora_r: 32
-lora_alpha: 16
-lora_dropout: 0.05
-lora_target_linear: true
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 2
-num_epochs: 1
-optimizer: adamw_bnb_8bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: false
-
-gradient_checkpointing: true
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_ratio: 0.1
-evals_per_epoch: 1
-saves_per_epoch: 1
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/streaming/README.md

@@ -1,50 +0,0 @@
-# Streaming Dataset Examples
-
-This directory contains example configurations for using Axolotl's streaming dataset
-functionality, which enables memory-efficient training with large datasets.
-
-## Examples
-
-Run the following examples with e.g. `axolotl train examples/streaming/sft.yaml`; no
-`axolotl preprocess` required!
-
-### Pretraining (`pretrain.yaml`)
-
-Demonstrates streaming configuration for pretraining tasks using the fineweb-edu dataset
-with SmolLM2-135M.
-
-- Uses `pretraining_dataset` configuration for automatic streaming
-- Multipack attention control to prevent cross-attention between packed sequences
-- Buffer size configuration for memory management
-
-### SFT (`sft.yaml`)
-
-Shows how to use streaming for supervised fine-tuning with the Alpaca dataset.
-
-- Explicit `streaming: true` flag for SFT datasets
-- Memory-efficient training on instruction datasets
-- Evaluation datasets are currently not streamed
-
-## Key Configuration Options
-
-### `streaming`
-- Enables streaming mode for standard datasets
-- Automatically enabled for `pretraining_dataset`
-
-### `streaming_multipack_buffer_size`
-- Controls buffer size for sample packing (default: 10,000)
-- Larger values improve packing efficiency but use more memory
-- Adjust based on available memory
-
-### `shuffle_merged_datasets`
-- Enables shuffling of streaming datasets
-- Requires additional memory for shuffle buffer
-
-### `sample_packing`
-- Packs multiple samples into single sequences
-- Minimize per-step padding tokens
-
-## Performance Tips
-
-- Download small / frequently-used datasets locally for better performance
-- Larger buffer sizes improve packing efficiency

examples/streaming/pretrain.yaml

@@ -1,57 +0,0 @@
-base_model: HuggingFaceTB/SmolLM2-135M
-
-# Streaming pretraining configuration
-pretraining_dataset:
-  - path: HuggingFaceFW/fineweb-edu
-    name: sample-10BT
-    type: pretrain
-    text_column: text
-    split: train
-
-# Streaming-specific settings
-streaming_multipack_buffer_size: 10000
-shuffle_merged_datasets: true
-
-# Training configuration
-max_steps: 1000
-output_dir: ./outputs/smollm2-135m-pretrain-streaming
-
-# Sequence and packing settings
-sequence_len: 1024
-sample_packing: true
-pretrain_multipack_attn: true  # Prevent cross-attention between packed sequences
-flash_attention: true
-
-# Batch size settings
-gradient_accumulation_steps: 8
-micro_batch_size: 1
-
-# Optimizer and scheduler
-optimizer: adamw_torch
-lr_scheduler: cosine
-learning_rate: 5e-4
-warmup_ratio: 0.1
-weight_decay: 0.01
-
-# Precision and performance
-bf16: auto
-tf32: true
-
-# Logging and checkpointing
-logging_steps: 10
-save_strategy: steps
-save_steps: 250
-save_total_limit: 3
-
-# Weights & Biases (optional)
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-# Special tokens
-special_tokens:
-  pad_token: "<|endoftext|>"
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/streaming/sft.yaml

@@ -1,55 +0,0 @@
-base_model: HuggingFaceTB/SmolLM2-135M
-
-# Dataset configuration
-datasets:
-  - path: tatsu-lab/alpaca
-    type: alpaca
-    split: train
-
-# Streaming-specific settings
-streaming: true
-streaming_multipack_buffer_size: 10000
-shuffle_merged_datasets: true
-
-# Training configuration
-max_steps: 1000
-output_dir: ./outputs/smollm2-135m-sft-streaming
-
-# Sequence and packing settings
-sequence_len: 1024
-sample_packing: true
-flash_attention: true
-
-# Batch size settings
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-
-# Optimizer and scheduler
-optimizer: adamw_torch
-lr_scheduler: cosine
-learning_rate: 2e-4
-warmup_ratio: 0.1
-weight_decay: 0.0
-
-# Precision and performance
-bf16: auto
-tf32: true
-
-# Logging and checkpointing
-logging_steps: 10
-save_strategy: steps
-save_steps: 100
-save_total_limit: 3
-
-# Weights & Biases (optional)
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-# Special tokens
-special_tokens:
-  pad_token: "<|endoftext|>"
-
-# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/voxtral/README.md

@@ -22,19 +22,9 @@ pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
 # audio
 pip3 install librosa==0.11.0
 pip3 install 'mistral_common[audio]==1.8.3'
-
-# Install CCE https://docs.axolotl.ai/docs/custom_integrations.html#cut-cross-entropy
-python scripts/cutcrossentropy_install.py | sh
 ```
 
-3. Download sample dataset files
-
-```bash
-# for text + audio only
-wget https://huggingface.co/datasets/Nanobit/text-audio-2k-test/resolve/main/En-us-African_elephant.oga
-```
-
-4. Run the finetuning example:
+3. Run the finetuning example:
 
 ```bash
 # text only

pyproject.toml

@@ -32,7 +32,7 @@ line-length = 88
 target-version = "py310"
 
 [tool.ruff.lint]
-select = ["E", "F", "W", "C90", "B", "I"]
+select = ["E", "F", "W", "C90", "B"]
 ignore = [
     "E203",  # Whitespace before ':'
     "E501",  # Line too long

requirements.txt

@@ -13,12 +13,12 @@ packaging==23.2
 
 huggingface_hub>=0.33.0
 peft>=0.17.0
-transformers==4.56.1
+transformers==4.55.3
 tokenizers>=0.21.1
-accelerate==1.10.1
+accelerate==1.10.0
 datasets==4.0.0
 deepspeed>=0.17.0
-trl==0.23.0
+trl==0.21.0
 hf_xet==1.1.5
 kernels==0.9.0
 trackio
@@ -64,10 +64,10 @@ langdetect==1.0.9
 immutabledict==4.2.0
 antlr4-python3-runtime==4.13.2
 
-torchao==0.13.0
+torchao==0.12.0
 schedulefree==1.4.1
 
 axolotl-contribs-lgpl==0.0.6
 axolotl-contribs-mit==0.0.5
 
-mistral-common==1.8.5
+mistral-common==1.8.3

scripts/cutcrossentropy_install.py

@@ -29,5 +29,5 @@ UV_PREFIX = "uv " if USE_UV else ""
 
 print(
     UNINSTALL_PREFIX
-    + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c5aa3ef"'
+    + f'{UV_PREFIX}pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5"'
 )

setup.py

@@ -124,9 +124,10 @@ extras_require = {
     "ring-flash-attn": [
         "flash-attn==2.8.3",
         "ring-flash-attn>=0.1.7",
+        "yunchang==0.6.0",
     ],
     "deepspeed": [
-        "deepspeed==0.17.5",
+        "deepspeed==0.17.2",
         "deepspeed-kernels",
     ],
     "mamba-ssm": [
@@ -161,7 +162,6 @@ extras_require = {
     "llmcompressor": [
         "llmcompressor==0.5.1",
     ],
-    "fbgemm-gpu": ["fbgemm-gpu-genai>=1.2.0"],
 }
 install_requires, dependency_links, extras_require_build = parse_requirements(
     extras_require

src/axolotl/cli/__init__.py

@@ -4,7 +4,5 @@ import os
 
 from axolotl.logging_config import configure_logging
 
-os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
-os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")
-
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
 configure_logging()

src/axolotl/cli/args.py

@@ -14,13 +14,9 @@ class PreprocessCliArgs:
     prompter: Optional[str] = field(default=None)
     download: Optional[bool] = field(default=True)
     iterable: Optional[bool] = field(
-        default=False,
+        default=None,
         metadata={
-            "help": (
-                "Deprecated in v0.13.0, will be removed in v0.14.0. For streaming "
-                "datasets, use 'axolotl train' and set 'streaming: true' in your YAML "
-                "config, or pass --streaming instead in the CLI."
-            )
+            "help": "Use IterableDataset for streaming processing of large datasets"
         },
     )
 
@@ -115,7 +111,6 @@ class QuantizeCliArgs:
     quantize_embedding: Optional[bool] = field(default=None)
     group_size: Optional[int] = field(default=None)
     output_dir: Optional[str] = field(default=None)
-    hub_model_id: Optional[str] = field(default=None)
 
 
 @dataclass

src/axolotl/cli/cloud/__init__.py

@@ -67,8 +67,16 @@ def do_cli_lm_eval(
     cloud_config: Path | str,
     config: Path | str,
 ) -> None:
-    cloud_cfg = load_cloud_cfg(cloud_config)
-    cloud = ModalCloud(cloud_cfg)
+    cloud_cfg: DictDefault = load_cloud_cfg(cloud_config)
+    provider = cloud_cfg.provider or "modal"
+    cloud: Cloud | None
+    if provider == "modal":
+        cloud = ModalCloud(cloud_cfg)
+    elif provider == "baseten":
+        cloud = BasetenCloud(cloud_cfg.to_dict())
+    else:
+        raise ValueError(f"Unsupported cloud provider: {provider}")
+
     with open(config, "r", encoding="utf-8") as file:
         config_yaml = file.read()
     cloud.lm_eval(config_yaml)

src/axolotl/cli/cloud/baseten/__init__.py

@@ -46,3 +46,23 @@ class BasetenCloud(Cloud):
             subprocess.run(  # nosec B603 B607
                 ["truss", "train", "push", "train_sft.py"], cwd=tmp_dir, check=False
             )
+
+    def lm_eval(
+        self,
+        config_yaml: str,
+    ):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            config = self.config.copy()
+            with open(tmp_dir + "/cloud.yaml", "w", encoding="utf-8") as cloud_fout:
+                yaml.dump(config, cloud_fout)
+            with open(tmp_dir + "/eval.yaml", "w", encoding="utf-8") as config_fout:
+                config_fout.write(config_yaml)
+            shutil.copyfile(
+                dirname(__file__) + "/template/eval.sh", tmp_dir + "/eval.sh"
+            )
+            shutil.copyfile(
+                dirname(__file__) + "/template/eval_sft.py", tmp_dir + "/eval_sft.py"
+            )
+            subprocess.run(  # nosec B603 B607
+                ["truss", "train", "push", "eval_sft.py"], cwd=tmp_dir, check=False
+            )

src/axolotl/cli/cloud/baseten/template/eval.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+set -eux
+
+export NCCL_SOCKET_IFNAME="^docker0,lo"
+export NCCL_IB_DISABLE=0
+export NCCL_TIMEOUT=1800000
+
+axolotl lm-eval eval.yaml

src/axolotl/cli/cloud/baseten/template/eval_sft.py

@@ -0,0 +1,81 @@
+"""
+Baseten Training Script for Axolotl
+"""
+
+# pylint: skip-file
+import yaml
+from truss.base import truss_config
+
+# Import necessary classes from the Baseten Training SDK
+from truss_train import definitions
+
+cloud_config = yaml.safe_load(open("cloud.yaml", "r"))
+gpu = cloud_config.get("gpu", "h100")
+gpu_count = (
+    1  # int(cloud_config.get("gpu_count", 1)) # only single GPU supported at the moment
+)
+node_count = (
+    1  # int(cloud_config.get("node_count", 1)) # only single node support for lmeval
+)
+project_name = cloud_config.get("project_name", "axolotl-project") or "axolotl-project"
+secrets = cloud_config.get("secrets", [])
+# launcher = cloud_config.get("launcher", "accelerate")
+# launcher_args = cloud_config.get("launcher_args", [])
+script_name = "eval.sh"
+
+# launcher_args_str = ""
+# if launcher_args:
+#     launcher_args_str = "-- " + " ".join(launcher_args)
+
+# 1. Define a base image for your training job
+# must use torch 2.7.0 for vllm
+BASE_IMAGE = "axolotlai/axolotl:main-py3.11-cu126-2.7.1"
+
+# 2. Define the Runtime Environment for the Training Job
+# This includes start commands and environment variables.a
+# Secrets from the baseten workspace like API keys are referenced using
+# `SecretReference`.
+
+env_vars = {
+    # "AXOLOTL_LAUNCHER": launcher,
+    # "AXOLOTL_LAUNCHER_ARGS": launcher_args_str,
+}
+for secret_name in secrets:
+    env_vars[secret_name] = definitions.SecretReference(name=secret_name)
+
+training_runtime = definitions.Runtime(
+    start_commands=[  # Example: list of commands to run your training script
+        f"/bin/sh -c 'chmod +x ./{script_name} && ./{script_name}'"
+    ],
+    environment_variables=env_vars,
+    cache_config=definitions.CacheConfig(
+        enabled=True,
+    ),
+    checkpointing_config=definitions.CheckpointingConfig(
+        enabled=True,
+    ),
+)
+
+# 3. Define the Compute Resources for the Training Job
+training_compute = definitions.Compute(
+    node_count=node_count,
+    accelerator=truss_config.AcceleratorSpec(
+        accelerator=truss_config.Accelerator.H100,
+        count=gpu_count,
+    ),
+)
+
+# 4. Define the Training Job
+# This brings together the image, compute, and runtime configurations.
+my_training_job = definitions.TrainingJob(
+    image=definitions.Image(base_image=BASE_IMAGE),
+    compute=training_compute,
+    runtime=training_runtime,
+)
+
+
+# This config will be pushed using the Truss CLI.
+# The association of the job to the project happens at the time of push.
+first_project_with_job = definitions.TrainingProject(
+    name=project_name, job=my_training_job
+)

src/axolotl/cli/cloud/baseten/template/train_sft.py

@@ -44,6 +44,12 @@ training_runtime = definitions.Runtime(
         f"/bin/sh -c 'chmod +x ./{script_name} && ./{script_name}'"
     ],
     environment_variables=env_vars,
+    cache_config=definitions.CacheConfig(
+        enabled=True,
+    ),
+    checkpointing_config=definitions.CheckpointingConfig(
+        enabled=True,
+    ),
 )
 
 # 3. Define the Compute Resources for the Training Job

src/axolotl/cli/config.py

@@ -23,8 +23,7 @@ from axolotl.utils.config import (
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.logging import get_logger
 from axolotl.utils.mlflow_ import setup_mlflow_env_vars
-from axolotl.utils.tee import prepare_debug_log
-from axolotl.utils.trainer import prepare_optim_env
+from axolotl.utils.trainer import prepare_opinionated_env, prepare_optim_env
 from axolotl.utils.wandb_ import setup_wandb_env_vars
 
 LOG = get_logger(__name__)
@@ -228,11 +227,8 @@ def load_cfg(
         },
     )
 
-    # NOTE(djsaunde): We start outputting to output_dir/debug.log at this point since we
-    # have to wait for cfg.output to be resolved. We could call this earlier if we write
-    # to a temporary file, and then move it later.
-    prepare_debug_log(cfg)
     prepare_optim_env(cfg)
+    prepare_opinionated_env(cfg)
     normalize_config(cfg)
     normalize_cfg_datasets(cfg)
     setup_wandb_env_vars(cfg)
@@ -245,6 +241,7 @@ def load_cfg(
         for k, v in cfg.items()
         if v is not None
     }
+
     LOG.info(
         "config:\n%s",
         json.dumps(cfg_to_log, indent=2, default=str, sort_keys=True),

src/axolotl/cli/inference.py

@@ -14,12 +14,10 @@ from transformers import GenerationConfig, TextIteratorStreamer, TextStreamer
 from axolotl.cli.args import InferenceCliArgs
 from axolotl.cli.config import load_cfg
 from axolotl.cli.utils import load_model_and_tokenizer
-from axolotl.cli.utils.diffusion import (
-    diffusion_inference,
-    launch_diffusion_gradio_ui,
+from axolotl.utils.chat_templates import (
+    get_chat_template,
+    get_chat_template_from_config,
 )
-from axolotl.integrations.base import PluginManager
-from axolotl.utils.chat_templates import get_chat_template_from_config
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.logging import get_logger
 
@@ -34,7 +32,6 @@ def get_multi_line_input() -> str:
         Possibly multi-line, possibly empty stdin input as a string.
     """
     print("Give me an instruction (Ctrl + D to submit): ")
-    print("=" * 80)
 
     instruction = ""
     for line in sys.stdin:
@@ -49,9 +46,9 @@ def do_inference(
     cli_args: InferenceCliArgs,
 ):
     """
-    Runs inference on the command line in a loop. User input is accepted, a chat
-    template is (optionally) applied, and the model specified in the `axolotl` config is
-    used to generate completions according to a default generation config.
+    Runs inference on the command line in a loop. User input is accepted, a chat template
+    is (optionally) applied, and the model specified in the `axolotl` config is used to
+    generate completions according to a default generation config.
 
     Args:
         cfg: Dictionary mapping `axolotl` config keys to values.
@@ -67,31 +64,17 @@ def do_inference(
             importlib.import_module("axolotl.prompters"), prompter
         )
     elif cfg.chat_template:
-        chat_template_str = get_chat_template_from_config(
-            cfg, ds_cfg=None, tokenizer=tokenizer
-        )
-    elif cfg.datasets and cfg.datasets[0].type == "chat_template":
+        chat_template_str = get_chat_template(cfg.chat_template, tokenizer=tokenizer)
+    elif cfg.datasets[0].type == "chat_template":
         chat_template_str = get_chat_template_from_config(
             cfg=cfg, ds_cfg=cfg.datasets[0], tokenizer=tokenizer
         )
 
     model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
-    # Detect diffusion mode
-    plugin_manager = PluginManager.get_instance()
-    is_diffusion = any(
-        plugin.__class__.__name__ == "DiffusionPlugin"
-        for plugin in plugin_manager.plugins.values()
-    )
-
-    if is_diffusion:
-        print("=" * 80)
-        print("Commands:")
-        print(":complete N -> completion mode with N tokens (default 64)")
-        print(":mask R     -> random masking with ratio R (0.0–1.0)")
-
     while True:
         print("=" * 80)
+        # support for multiline inputs
         instruction = get_multi_line_input()
         if not instruction:
             return
@@ -121,19 +104,9 @@ def do_inference(
         else:
             batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
 
-        print("=" * 80)
+        print("=" * 40)
         model.eval()
         with torch.no_grad():
-            if is_diffusion:
-                diffusion_inference(
-                    model=model,
-                    tokenizer=tokenizer,
-                    cfg=cfg,
-                    prompt=prompt,
-                    chat_template_str=chat_template_str,
-                )
-                continue
-
             generation_config = GenerationConfig(
                 repetition_penalty=1.1,
                 max_new_tokens=1024,
@@ -156,7 +129,7 @@ def do_inference(
                 generation_config=generation_config,
                 streamer=streamer,
             )
-        print("=" * 80)
+        print("=" * 40)
         print(tokenizer.decode(generated["sequences"].cpu().tolist()[0]))
 
 
@@ -186,33 +159,10 @@ def do_inference_gradio(
             importlib.import_module("axolotl.prompters"), prompter
         )
     elif cfg.chat_template:
-        chat_template_str = get_chat_template_from_config(
-            cfg, ds_cfg=None, tokenizer=tokenizer
-        )
-    elif cfg.datasets and cfg.datasets[0].type == "chat_template":
-        chat_template_str = get_chat_template_from_config(
-            cfg=cfg, ds_cfg=cfg.datasets[0], tokenizer=tokenizer
-        )
+        chat_template_str = get_chat_template(cfg.chat_template, tokenizer=tokenizer)
 
     model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
-    # Detect diffusion mode
-    plugin_manager = PluginManager.get_instance()
-    is_diffusion = any(
-        plugin.__class__.__name__ == "DiffusionPlugin"
-        for plugin in plugin_manager.plugins.values()
-    )
-
-    if is_diffusion:
-        launch_diffusion_gradio_ui(
-            model=model,
-            tokenizer=tokenizer,
-            cfg=cfg,
-            prompter_module=prompter_module,
-            chat_template_str=chat_template_str,
-        )
-        return
-
     def generate(instruction):
         if not instruction:
             return

src/axolotl/cli/main.py

@@ -26,7 +26,7 @@ from axolotl.cli.utils import (
     launch_training,
 )
 from axolotl.integrations.lm_eval.cli import lm_eval
-from axolotl.utils import set_pytorch_cuda_alloc_conf
+from axolotl.utils import patch_optimized_env
 from axolotl.utils.logging import get_logger
 from axolotl.utils.schemas.config import AxolotlInputConfig
 
@@ -44,7 +44,7 @@ def cli():
     """Axolotl CLI - Train and fine-tune large language models"""
     print_axolotl_text_art()
     load_dotenv()
-    set_pytorch_cuda_alloc_conf()
+    patch_optimized_env()
 
 
 @cli.command()

src/axolotl/cli/preprocess.py

@@ -35,20 +35,10 @@ def do_preprocess(cfg: DictDefault, cli_args: PreprocessCliArgs) -> None:
     check_accelerate_default_config()
     check_user_token()
 
-    if cli_args.iterable:
-        LOG.error(
-            "The --iterable CLI argument for 'axolotl preprocess' is no longer "
-            "supported. For training, set 'streaming: true' in your YAML config or "
-            "pass '--streaming' in your 'axolotl train' command for on-the-fly "
-            "preprocessing."
-        )
-        return
-
     for key in ["skip_prepare_dataset", "pretraining_dataset"]:
         if cfg.get(key):
             LOG.error(
-                f"You have set `{key}:`. `preprocess` is not needed. Run the 'axolotl "
-                "train' CLI directly instead."
+                f"You have set `{key}:`. `preprocess` is not needed. Run the `axolotl train` CLI directly instead."
             )
             return
 

src/axolotl/cli/quantize.py

@@ -5,17 +5,12 @@ CLI to post-training quantize a model using torchao
 from pathlib import Path
 from typing import Union
 
-from transformers import AutoConfig, AutoModelForCausalLM, TorchAoConfig
+from transformers import AutoModelForCausalLM
 
 from axolotl.cli.config import load_cfg
 from axolotl.loaders import load_tokenizer
 from axolotl.utils.logging import get_logger
-from axolotl.utils.quantization import (
-    TorchAOQuantDType,
-    get_quantization_config,
-    quantization_config_to_str,
-    quantize_model,
-)
+from axolotl.utils.quantization import TorchIntDType, quantize_model_for_ptq
 
 LOG = get_logger(__name__)
 
@@ -48,13 +43,13 @@ def do_quantize(
             "No quantization configuration found. Please specify either qat or quantization in your config file."
         )
 
-    model_path = cli_args.get("base_model") or cfg.output_dir
+    model_path = cli_args.get("model_path") or cfg.output_dir
     if weight_dtype := cli_args.get("weight_dtype"):
-        weight_dtype = TorchAOQuantDType.from_string(weight_dtype)
+        weight_dtype = TorchIntDType[weight_dtype]
     else:
         weight_dtype = quantize_cfg.weight_dtype
     if activation_dtype := cli_args.get("activation_dtype"):
-        activation_dtype = TorchAOQuantDType.from_string(activation_dtype)
+        activation_dtype = TorchIntDType[activation_dtype]
     else:
         activation_dtype = quantize_cfg.activation_dtype
     group_size = cli_args.get("group_size") or quantize_cfg.group_size
@@ -62,15 +57,10 @@ def do_quantize(
         cli_args.get("quantize_embedding") or quantize_cfg.quantize_embedding
     )
     output_dir = cli_args.get("output_dir") or cfg.output_dir
-    hub_model_id = cli_args.get("hub_model_id") or cfg.hub_model_id
 
-    LOG.info(f"Loading model from {model_path}.")
+    LOG.info(f"Loading model from {model_path}...")
     tokenizer = load_tokenizer(cfg)
-    config = AutoConfig.from_pretrained(model_path)
-    torch_dtype = config.torch_dtype if hasattr(config, "torch_dtype") else None
-    model = AutoModelForCausalLM.from_pretrained(
-        model_path, device_map="auto", torch_dtype=torch_dtype
-    )
+    model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto")
 
     LOG.info(
         f"Quantizing model with configuration: \n"
@@ -80,21 +70,11 @@ def do_quantize(
         f"\tquantize_embedding: {quantize_embedding}"
     )
 
-    quantize_model(
+    quantize_model_for_ptq(
         model, weight_dtype, group_size, activation_dtype, quantize_embedding
     )
 
-    quantization_config = get_quantization_config(
-        weight_dtype, activation_dtype, group_size
-    )
-
-    ao_config = TorchAoConfig(
-        quant_type=quantization_config,
-        include_input_output_embeddings=quantize_embedding,
-    )
-    model.config.quantization_config = ao_config
-
-    LOG.info(f"Saving quantized model to: {str(Path(output_dir) / 'quantized')}.")
+    LOG.info(f"Saving quantized model to: {str(Path(output_dir) / 'quantized')}...")
     model.save_pretrained(
         str(Path(output_dir) / "quantized"),
         safe_serialization=False,
@@ -106,14 +86,4 @@ def do_quantize(
         progressbar=True,
         save_jinja_files=cfg.tokenizer_save_jinja_files,
     )
-
-    if hub_model_id:
-        hub_model_id = (
-            hub_model_id.rstrip("-")
-            + f"-{quantization_config_to_str[type(quantization_config)]}"
-        )
-        model.push_to_hub(hub_model_id, safe_serialization=False)
-        tokenizer.push_to_hub(hub_model_id)
-        LOG.info(f"Quantized model pushed to: {hub_model_id}.")
-
-    LOG.info(f"Quantized model saved to: {str(Path(output_dir) / 'quantized')}.")
+    LOG.info(f"Quantized model saved to: {str(Path(output_dir) / 'quantized')}...")

src/axolotl/cli/train.py

@@ -17,7 +17,6 @@ from axolotl.integrations.base import PluginManager
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, resolve_dtype
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.trainer import prepare_optim_env
 
 
 def do_train(cfg: DictDefault, cli_args: TrainerCliArgs):
@@ -60,6 +59,7 @@ def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
         config: Path to `axolotl` config YAML file.
         kwargs: Additional keyword arguments to override config file values.
     """
+
     parsed_cfg = load_cfg(config, **kwargs)
     parser = HfArgumentParser(TrainerCliArgs)
     parsed_cli_args, _ = parser.parse_args_into_dataclasses(
@@ -92,7 +92,6 @@ def ray_train_func(kwargs: dict):
     # cast `cfg` back to DictDefault (ray tune deepcopy has issues with DictDefault so needed it to be dict)
     # also renormalize the config now that TorchTrainer has spawned distributed workers
     cfg = DictDefault(kwargs["cfg"])
-    prepare_optim_env(cfg)
     normalize_config(cfg)
 
     # now that we are on the worker node, we can check `is_torch_bf16_gpu_available` to resolve dtype

src/axolotl/cli/utils/diffusion.py

@@ -1,374 +0,0 @@
-"""Helpers for diffusion-mode inference in CLI and Gradio."""
-
-from __future__ import annotations
-
-import gradio as gr
-from colorama import Fore, Style
-
-from axolotl.integrations.diffusion import generate, resolve_mask_token_id
-from axolotl.utils.dict import DictDefault
-
-
-def diffusion_inference(
-    model,
-    tokenizer,
-    cfg,
-    prompt: str,
-    chat_template_str: str | None = None,
-):
-    """Diffusion inference helper method."""
-    mode = "random"
-    completion_tokens = 0
-    target_mask_ratio = None
-    mode, completion_tokens, target_mask_ratio, cleaned = _parse_commands(prompt)
-
-    if cleaned:
-        prompt = cleaned
-
-    info = run_diffusion(
-        model=model,
-        tokenizer=tokenizer,
-        cfg=cfg,
-        prompt=prompt,
-        chat_template_str=chat_template_str,
-        mode=mode,
-        target_mask_ratio=target_mask_ratio,
-        completion_tokens=completion_tokens,
-    )
-    masked_text = info["masked_text"]
-    mask_ratio = info["mask_ratio"]
-    generated_ids = info["generated_ids"]
-    masked_positions = info["masked_positions"]
-    orig_ids = info["orig_ids"]
-
-    # Display with masked preview and colored diff
-    if masked_text is not None and mask_ratio is not None:
-        print(f"Masked ({mask_ratio:.1%}):\n{masked_text}\n")
-    if generated_ids is not None:
-        # Compute per-token style
-        styles: list[str] = []
-        for i, tid in enumerate(generated_ids):
-            if i in masked_positions:
-                if i < len(orig_ids) and tid == orig_ids[i]:
-                    styles.append("green")  # correct fill
-                elif i < len(orig_ids):
-                    styles.append("red")  # incorrect fill
-                else:
-                    styles.append("normal")  # appended
-            else:
-                same = i < len(orig_ids) and tid == orig_ids[i]
-                styles.append("dim" if same else "normal")
-
-        # Group contiguous spans by style
-        styled_spans: list[tuple[str, int, int]] = []
-        if generated_ids:
-            current_style = styles[0]
-            start = 0
-            for i in range(1, len(generated_ids)):
-                s = styles[i]
-                if s != current_style:
-                    styled_spans.append((current_style, start, i))
-                    current_style, start = s, i
-            styled_spans.append((current_style, start, len(generated_ids)))
-
-        out_parts = []
-        for style_name, a, b in styled_spans:
-            chunk_text = tokenizer.decode(generated_ids[a:b], skip_special_tokens=False)
-            if style_name == "green":
-                out_parts.append(Fore.GREEN + chunk_text + Style.RESET_ALL)
-            elif style_name == "red":
-                out_parts.append(Fore.RED + chunk_text + Style.RESET_ALL)
-            else:
-                if style_name == "dim":
-                    out_parts.append(Style.DIM + chunk_text + Style.RESET_ALL)
-                else:
-                    out_parts.append(chunk_text)
-        print("Generated:\n" + "".join(out_parts))
-    else:
-        print("Generated:\n(no output)")
-
-
-def _parse_commands(text: str):
-    """
-    Parse leading diffusion commands.
-
-    Supported at start of input (can be chained):
-      :complete N  -> completion mode with N tokens (default 64)
-      :mask R      -> random masking with ratio R in [0, 1]
-    """
-    tokens = text.strip().split()
-    i = 0
-    mode = "random"
-    completion_tokens = 0
-    target_mask_ratio = None
-    consumed = 0
-    while i < len(tokens) and tokens[i].startswith(":"):
-        cmd = tokens[i]
-        i += 1
-        consumed = i
-        if cmd == ":complete":
-            mode = "completion"
-            if i < len(tokens):
-                try:
-                    completion_tokens = int(tokens[i])
-                    i += 1
-                    consumed = i
-                except Exception:
-                    completion_tokens = 64
-            else:
-                completion_tokens = 64
-        elif cmd == ":mask":
-            mode = "random"
-            if i < len(tokens):
-                try:
-                    target_mask_ratio = float(tokens[i])
-                    i += 1
-                    consumed = i
-                except Exception:
-                    target_mask_ratio = None
-        else:
-            i -= 1
-            consumed = i
-            break
-
-    cleaned = " ".join(tokens[consumed:])
-
-    return mode, completion_tokens, target_mask_ratio, cleaned
-
-
-def run_diffusion(
-    *,
-    model,
-    tokenizer,
-    cfg: DictDefault,
-    prompt: str,
-    chat_template_str: str | None,
-    mode: str = "random",
-    target_mask_ratio: float | None = None,
-    completion_tokens: int = 0,
-):
-    """Run a single diffusion generation and return a structured result dict."""
-    if chat_template_str:
-        batch = tokenizer.apply_chat_template(
-            [{"role": "user", "content": prompt}],
-            return_tensors="pt",
-            add_special_tokens=True,
-            add_generation_prompt=True,
-            chat_template=chat_template_str,
-            tokenize=True,
-            return_dict=True,
-        )
-    else:
-        batch = tokenizer(prompt, return_tensors="pt", add_special_tokens=True)
-
-    mask_token_id = resolve_mask_token_id(tokenizer, cfg, allow_add=False)
-
-    seq = batch["input_ids"].to(cfg.device)
-    gen_mode = "completion" if mode == "completion" else "random"
-    comp_tokens = int(completion_tokens) if gen_mode == "completion" else 0
-
-    result = generate(
-        model,
-        tokenizer,
-        original_sequence=seq[:1],
-        num_diffusion_steps=cfg.diffusion.num_diffusion_steps,
-        temperature=cfg.diffusion.generation_temperature,
-        mask_token_id=int(mask_token_id),
-        mode=gen_mode,  # type: ignore[arg-type]
-        completion_tokens=comp_tokens,
-        target_mask_ratio=target_mask_ratio,
-    )
-
-    masked_text = result.get("masked") if isinstance(result, dict) else None
-    mask_ratio = result.get("mask_ratio") if isinstance(result, dict) else None
-    generated_ids = result.get("generated_ids") if isinstance(result, dict) else None
-    masked_positions = (
-        set(result.get("masked_positions") or []) if isinstance(result, dict) else set()
-    )
-    orig_ids = seq[0].detach().cpu().tolist()
-
-    return {
-        "masked_text": masked_text,
-        "mask_ratio": mask_ratio,
-        "generated_ids": generated_ids,
-        "masked_positions": masked_positions,
-        "orig_ids": orig_ids,
-    }
-
-
-def render_html(
-    *,
-    generated_ids: list[int] | None,
-    orig_ids: list[int],
-    masked_positions: set[int],
-    tokenizer,
-) -> str:
-    """Render HTML visualizing diffusion outputs."""
-    if not generated_ids:
-        return "<pre>Generated:\n(no output)</pre>"
-
-    def _style_for(i: int, tid: int) -> str:
-        if i in masked_positions:
-            if i < len(orig_ids) and tid == orig_ids[i]:
-                return "green"
-            if i < len(orig_ids):
-                return "red"
-            return "normal"
-        same = i < len(orig_ids) and tid == orig_ids[i]
-        return "dim" if same else "normal"
-
-    # Group contiguous spans by style to reduce HTML size
-    spans: list[tuple[str, int, int]] = []
-    if generated_ids:
-        cur = _style_for(0, generated_ids[0])
-        start = 0
-        for i in range(1, len(generated_ids)):
-            s = _style_for(i, generated_ids[i])
-            if s != cur:
-                spans.append((cur, start, i))
-                cur, start = s, i
-        spans.append((cur, start, len(generated_ids)))
-
-    html_parts = []
-    for style_name, a, b in spans:
-        txt = tokenizer.decode(generated_ids[a:b], skip_special_tokens=False)
-        if style_name == "green":
-            html_parts.append(f'<span style="color:#2e7d32">{txt}</span>')
-        elif style_name == "red":
-            html_parts.append(f'<span style="color:#c62828">{txt}</span>')
-        elif style_name == "dim":
-            html_parts.append(f'<span style="opacity:0.6">{txt}</span>')
-        else:
-            html_parts.append(txt)
-
-    legend = (
-        '<div style="font-size:0.9em;margin-bottom:4px">'
-        '<span style="color:#2e7d32">correct</span>, '
-        '<span style="color:#c62828">incorrect</span>, '
-        '<span style="opacity:0.6">unchanged</span>'
-        "</div>"
-    )
-
-    return (
-        legend
-        + '<pre style="white-space:pre-wrap">Generated:\n'
-        + "".join(html_parts)
-        + "</pre>"
-    )
-
-
-def launch_diffusion_gradio_ui(
-    *,
-    model,
-    tokenizer,
-    cfg: DictDefault,
-    prompter_module=None,
-    chat_template_str: str | None = None,
-):
-    """Build and launch a simple Gradio UI for diffusion inference."""
-    with gr.Blocks(
-        title=cfg.get("gradio_title", "Axolotl Diffusion Interface")
-    ) as demo:
-        gr.Markdown(
-            """
-            ## Axolotl Diffusion Inference
-            - Mode "Random" masks tokens at a target ratio and fills them.
-            - Mode "Completion" appends N masked tokens at the end and fills them.
-            """
-        )
-
-        with gr.Row():
-            mode = gr.Radio(
-                choices=["random", "completion"],
-                value="random",
-                label="Mode",
-            )
-            mask_ratio = gr.Slider(
-                minimum=0.0,
-                maximum=1.0,
-                step=0.05,
-                value=0.4,
-                label="Mask ratio (random mode)",
-                interactive=True,
-            )
-            completion_tokens = gr.Number(
-                value=64,
-                precision=0,
-                label="Completion tokens (completion mode)",
-                interactive=True,
-                visible=False,
-            )
-
-        instruction = gr.Textbox(label="Instruction", lines=6)
-        run_btn = gr.Button("Generate")
-
-        masked_preview = gr.Textbox(label="Masked preview", lines=6)
-        html_out = gr.HTML(label="Generated")
-
-        def _toggle_controls(selected_mode: str):
-            return (
-                gr.update(visible=(selected_mode == "random")),
-                gr.update(visible=(selected_mode == "completion")),
-            )
-
-        mode.change(
-            _toggle_controls,
-            inputs=[mode],
-            outputs=[mask_ratio, completion_tokens],
-        )
-
-        def _gen(instruction_text: str, selected_mode: str, mratio: float, ctoks: int):
-            if not instruction_text:
-                return "", "<pre>Generated:\n(no output)</pre>"
-
-            if prompter_module:
-                prompt: str = next(
-                    prompter_module().build_prompt(
-                        instruction=instruction_text.strip("\n")
-                    )
-                )
-            else:
-                prompt = instruction_text.strip()
-
-            info = run_diffusion(
-                model=model,
-                tokenizer=tokenizer,
-                cfg=cfg,
-                prompt=prompt,
-                chat_template_str=chat_template_str,
-                mode=selected_mode,
-                target_mask_ratio=mratio if selected_mode == "random" else None,
-                completion_tokens=int(ctoks) if selected_mode == "completion" else 0,
-            )
-
-            masked_text = info.get("masked_text")
-            mask_ratio_val = info.get("mask_ratio")
-            generated_ids = info.get("generated_ids")
-            masked_positions = info.get("masked_positions") or set()
-            orig_ids = info.get("orig_ids") or []
-
-            preview = (
-                f"Masked ({mask_ratio_val:.1%}):\n{masked_text}"
-                if masked_text is not None and mask_ratio_val is not None
-                else ""
-            )
-            html = render_html(
-                generated_ids=generated_ids,
-                orig_ids=orig_ids,
-                masked_positions=masked_positions,
-                tokenizer=tokenizer,
-            )
-            return preview, html
-
-        run_btn.click(
-            _gen,
-            inputs=[instruction, mode, mask_ratio, completion_tokens],
-            outputs=[masked_preview, html_out],
-        )
-
-        demo.queue().launch(
-            show_api=False,
-            share=cfg.get("gradio_share", True),
-            server_name=cfg.get("gradio_server_name", "127.0.0.1"),
-            server_port=cfg.get("gradio_server_port", None),
-        )

src/axolotl/common/datasets.py

@@ -55,11 +55,13 @@ def load_datasets(
     """
     tokenizer = load_tokenizer(cfg)
     processor = load_processor(cfg, tokenizer=tokenizer) if cfg.processor_type else None
+    preprocess_iterable = getattr(cli_args, "iterable", False)
 
     train_dataset, eval_dataset, total_num_steps, prompters = prepare_datasets(
         cfg,
         tokenizer,
         processor=processor,
+        preprocess_iterable=preprocess_iterable,
     )
 
     if (

src/axolotl/core/builders/base.py

@@ -36,6 +36,7 @@ from axolotl.utils.callbacks import (
     SaveModelOnFirstStepCallback,
 )
 from axolotl.utils.callbacks.profiler import PytorchProfilerCallback
+from axolotl.utils.callbacks.tokens_per_second import TokensPerSecondCallback
 from axolotl.utils.distributed import build_parallelism_config
 from axolotl.utils.schemas.enums import CustomSupportedOptimizers
 
@@ -144,6 +145,12 @@ class TrainerBuilderBase(abc.ABC):
                     profiler_steps_start=self.cfg.profiler_steps_start,
                 )
             )
+        if self.cfg.include_tkps:
+            callbacks.append(
+                TokensPerSecondCallback(
+                    self.cfg.tensor_parallel_size, self.cfg.context_parallel_size
+                )
+            )
 
         return callbacks
 
@@ -435,7 +442,7 @@ class TrainerBuilderBase(abc.ABC):
             # don't use the HF gradient checkpointing, manually wrap
             training_args_kwargs["gradient_checkpointing"] = False
             training_args_kwargs["activation_offloading"] = True
-        elif self.cfg.gradient_checkpointing is not None:
+        elif self.cfg.gradient_checkpointing:
             training_args_kwargs["gradient_checkpointing"] = (
                 self.cfg.gradient_checkpointing
             )

src/axolotl/core/builders/causal.py

@@ -10,7 +10,6 @@ import transformers
 from transformers import (
     DataCollatorWithFlattening,
     EarlyStoppingCallback,
-    Trainer,
 )
 from trl.trainer.utils import RewardDataCollatorWithPadding
 
@@ -36,7 +35,6 @@ from axolotl.utils.callbacks import (
 )
 from axolotl.utils.callbacks.lisa import lisa_callback_factory
 from axolotl.utils.callbacks.qat import QATCallback
-from axolotl.utils.callbacks.tokens_per_second import TokensPerSecondCallback
 from axolotl.utils.chat_templates import get_chat_template_from_config
 from axolotl.utils.collators import (
     BatchSamplerDataCollatorForSeq2Seq,
@@ -76,12 +74,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         if self.cfg.qat:
             callbacks.append(QATCallback(self.cfg.qat))
 
-        if self.cfg.include_tkps:
-            callbacks.append(
-                TokensPerSecondCallback(
-                    self.cfg.tensor_parallel_size, self.cfg.context_parallel_size
-                )
-            )
         return callbacks
 
     def get_post_trainer_create_callbacks(self, trainer):
@@ -348,10 +340,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
 
         if self.cfg.reward_model:
             training_args_cls = AxolotlRewardConfig
-            if self.cfg.center_rewards_coefficient is not None:
-                training_arguments_kwargs["center_rewards_coefficient"] = (
-                    self.cfg.center_rewards_coefficient
-                )
         elif self.cfg.process_reward_model:
             training_args_cls = AxolotlPRMConfig
         else:
@@ -395,11 +383,10 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 **data_collator_kwargs,
             )
         sig = inspect.signature(trainer_cls)
-        if "processing_class" in sig.parameters or issubclass(trainer_cls, Trainer):
+        if "processing_class" in sig.parameters:
             trainer_kwargs["processing_class"] = self.tokenizer
         elif "tokenizer" in sig.parameters:
             trainer_kwargs["tokenizer"] = self.tokenizer
-
         if (
             trainer_cls not in [AxolotlRewardTrainer, AxolotlPRMTrainer]
             and self.cfg.datasets is not None

src/axolotl/core/builders/rl.py

@@ -120,11 +120,6 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
         if self.cfg.use_wandb:
             training_args_kwargs["run_name"] = self.cfg.wandb_name
 
-        if self.cfg.max_prompt_len:
-            training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
-        else:
-            training_args_kwargs["max_prompt_length"] = self.cfg.sequence_len
-
         training_args_cls = None
         blocklist_args_kwargs = []
         if self.cfg.rl is RLType.SIMPO:
@@ -134,16 +129,10 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
             if self.cfg.cpo_alpha is not None:
                 training_args_kwargs["cpo_alpha"] = self.cfg.cpo_alpha
 
-            # Handle when max_prompt_length == max_length from defaults
-            # CPOTrainer requires strictly less than
-            if (
-                training_args_kwargs["max_prompt_length"]
-                == training_args_kwargs["max_length"]
-            ):
-                training_args_kwargs["max_prompt_length"] -= 1
-
         elif self.cfg.rl is RLType.ORPO:
             training_args_cls = AxolotlORPOConfig
+            if self.cfg.max_prompt_len:
+                training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
 
         elif self.cfg.rl is RLType.KTO:
             training_args_cls = AxolotlKTOConfig
@@ -155,6 +144,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
                 self.cfg.kto_undesirable_weight or 1.0
             )
 
+            if self.cfg.max_prompt_len:
+                training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
+
         elif self.cfg.rl is RLType.GRPO:
             training_args_cls = GRPOStrategy.get_training_args_class()
             training_args_kwargs.update(GRPOStrategy.set_training_args_kwargs(self.cfg))

src/axolotl/core/datasets/transforms/chat_builder.py

@@ -8,7 +8,7 @@ from typing import Any, Mapping
 
 def chat_message_transform_builder(
     train_on_inputs=False,
-    conversations_field: str = "messages",
+    conversations_field: str = "conversations",
     message_field_role: str | list[str] | None = None,  # commonly "role"
     message_field_content: str | list[str] | None = None,  # commonly "content"
     message_field_training: str | list[str] | None = None,  # commonly "weight"
@@ -20,13 +20,13 @@ def chat_message_transform_builder(
             If True, the transform will train on the inputs. If False, the transform will train on the targets.
             Defaults to False.
         conversations_field (str, optional):
-            The field name of the conversations. Defaults to "messages".
+            The field name of the conversations. Defaults to "conversations".
         message_field_role (str | list[str], optional):
-            The field name of the role.
+            The field name of the role. Defaults to "role".
         message_field_content (str | list[str], optional):
-            The field name of the message content.
+            The field name of the message content. Defaults to "content".
         message_field_training (str | list[str], optional):
-            The field name of the train/weight.
+            The field name of the train/weight. Defaults to "weight".
 
     Returns:
         Callable:

src/axolotl/core/trainers/base.py

@@ -49,13 +49,6 @@ from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 
 LOG = get_logger(__name__)
 
-REDUCTION_FNS = {
-    "mean": torch.mean,
-    "min": torch.min,
-    "max": torch.max,
-    "sum": torch.sum,
-}
-
 
 class AxolotlTrainer(
     PackingMixin,
@@ -96,9 +89,7 @@ class AxolotlTrainer(
 
         super().__init__(*_args, **kwargs)
         self.train_data_collator = self.data_collator
-        self._stored_metrics = defaultdict(
-            lambda: defaultdict(lambda: {"values": [], "reduction": "mean"})
-        )
+        self._stored_metrics = defaultdict(lambda: defaultdict(list))
         if self.args.orpo_alpha:
             self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
 
@@ -351,10 +342,10 @@ class AxolotlTrainer(
             inputs_key = "labels" if "labels" in inputs else "input_ids"
             if hasattr(self.state, "num_tokens"):
                 self.state.num_tokens = (
-                    self.state.num_tokens + (inputs[inputs_key] != -100).sum().cpu()
+                    self.state.num_tokens + (inputs[inputs_key] != -100).sum()
                 )
             else:
-                self.state.num_tokens = (inputs[inputs_key] != -100).sum().cpu()
+                self.state.num_tokens = (inputs[inputs_key] != -100).sum()
 
         if self.args.orpo_alpha:
             return self.orpo_compute_loss(
@@ -371,11 +362,6 @@ class AxolotlTrainer(
             num_items_in_batch=num_items_in_batch,
         )
 
-    @override
-    def evaluate(self, *args, **kwargs):
-        LOG.info("Running evaluation step...")
-        return super().evaluate(*args, **kwargs)
-
     @staticmethod
     def orpo_concatenate_inputs(inputs, label_pad_token=-100, pad_token=0, device=None):
         concatenated_batch = {}
@@ -599,17 +585,9 @@ class AxolotlTrainer(
         """
         # logs either has 'loss' or 'eval_loss'
         train_eval = "train" if "loss" in logs else "eval"
-
-        for key, metric_data in self._stored_metrics[train_eval].items():
-            values = torch.tensor(metric_data["values"])  # type: ignore[arg-type]
-            reduction_type = metric_data["reduction"]
-
-            fn = REDUCTION_FNS.get(reduction_type)
-            if fn is None:
-                raise NotImplementedError(
-                    "Metric reduction must be one of [mean, min, max, sum]"
-                )
-            logs[key] = round(fn(values).item(), 4)
+        # Add averaged stored metrics to logs
+        for key, metrics in self._stored_metrics[train_eval].items():
+            logs[key] = torch.tensor(metrics).mean().item()
 
         if is_main_process():
             # Add memory usage
@@ -633,27 +611,10 @@ class AxolotlTrainer(
         return super().log(logs, start_time)
 
     def store_metrics(
-        self,
-        metrics: dict[str, float] | dict[str, tuple[int | float, str]],
-        train_eval: Literal["train", "eval"] = "train",
-        reduction: Literal["mean", "min", "max", "sum"] = "mean",
+        self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train"
     ) -> None:
-        """
-        Store metrics with specified reduction type.
-
-        Args:
-            metrics: Dictionary of metric names to values, or metric names to (value,
-                reduction_type) tuples.
-            train_eval: Whether this is for training or evaluation.
-        """
         for key, value in metrics.items():
-            if isinstance(value, tuple):
-                value, _reduction = value  # type: ignore[assignment]
-            else:
-                value, _reduction = value, reduction
-
-            self._stored_metrics[train_eval][key]["values"].append(value)
-            self._stored_metrics[train_eval][key]["reduction"] = _reduction
+            self._stored_metrics[train_eval][key].append(value)
 
     def _save_checkpoint(self, model, trial, **kwargs):
         # make sure the checkpoint dir exists, since trainer is flakey

src/axolotl/core/trainers/dpo/__init__.py

@@ -27,6 +27,7 @@ class DPOStrategy:
             training_args_kwargs["label_smoothing"] = cfg.dpo_label_smoothing
         training_args_kwargs["max_completion_length"] = None
         training_args_kwargs["max_length"] = cfg.sequence_len
+        training_args_kwargs["max_prompt_length"] = cfg.sequence_len
         training_args_kwargs["generate_during_eval"] = cfg.dpo_generate_during_eval
         if cfg.dpo_use_weighting is not None:
             training_args_kwargs["use_weighting"] = cfg.dpo_use_weighting

src/axolotl/datasets.py

@@ -1,17 +1,18 @@
-"""
-Module containing dataset functionality.
-
-We want this to be a wrapper for an existing dataset that we have loaded. Lets use the
-concept of middlewares to wrap each dataset. We'll use the collators later on to pad the
-datasets.
-"""
+"""Module containing Dataset functionality"""
 
+import torch
 from datasets import Dataset, IterableDataset
 
 from axolotl.utils.logging import get_logger
 
 from .prompt_tokenizers import PromptTokenizingStrategy
 
+# We want this to be a wrapper for an existing dataset that we have loaded
+# lets use the concept of middlewares to wrap each dataset, for example
+# ConstantLengthDataset(ShuffledDataset([TokenizedPromptDataset(alpaca_dataset)]))
+# let's check to ensure we don't truncate an item in the middle, we'll use
+# the collators later on to pad the datasets
+
 LOG = get_logger(__name__)
 
 
@@ -85,3 +86,133 @@ def wrap_dataset_for_tokenized_prompt(
             **map_kwargs,
         )
     return TokenizedPromptDataset(prompt_tokenizer, dataset, **kwargs)
+
+
+# TODO this isn't the best since it can't interleave datasets
+class ConstantLengthDataset(IterableDataset):
+    """Iterable dataset that returns constant length chunks of tokens from stream of
+    text files.
+
+    Args:
+        tokenizer: The processor used for processing the data.
+        dataset: Dataset with text files.
+        seq_length: Length of token sequences to return.
+    """
+
+    def __init__(
+        self,
+        tokenizer,
+        datasets,
+        seq_length=2048,
+    ):
+        self.tokenizer = tokenizer
+        self.concat_token_id = tokenizer.eos_token_id
+        self.datasets: list[IterableDataset] = datasets
+        self.seq_length = seq_length
+
+        vocab_size = len(tokenizer.get_vocab())
+
+        if vocab_size <= torch.iinfo(torch.int16).max:
+            self.tokens_dtype = torch.int16
+        elif vocab_size <= torch.iinfo(torch.int32).max:
+            self.tokens_dtype = torch.int32
+        else:
+            self.tokens_dtype = torch.int64
+
+    def __iter__(self):
+        buffer = {
+            "input_ids": [],
+            "attention_mask": [],
+            "labels": [],
+            "position_ids": [],
+        }
+        buffer_len = 0
+        for dataset in self.datasets:
+            idx = 0
+            iterator = iter(dataset)
+            more_examples = True
+            while more_examples:
+                try:
+                    example = next(iterator)
+                    idx += 1
+                except StopIteration:
+                    more_examples = False
+                    example = None
+
+                add_concat_token = False
+                if example:
+                    example_len = len(example["input_ids"])
+                    add_concat_token = example["input_ids"][-1] != self.concat_token_id
+                else:
+                    example_len = 0
+
+                if not example_len or (
+                    buffer_len + int(add_concat_token) + example_len > self.seq_length
+                ):
+                    if buffer["input_ids"]:
+                        input_ids = torch.cat(buffer["input_ids"], dim=-1)[
+                            : self.seq_length
+                        ]
+                        attention_mask = torch.cat(buffer["attention_mask"], dim=-1)[
+                            : self.seq_length
+                        ]
+                        position_ids = torch.cat(buffer["position_ids"], dim=-1)[
+                            : self.seq_length
+                        ]
+                        labels = torch.cat(buffer["labels"], dim=-1)[: self.seq_length]
+                        if labels.size() == input_ids.size() and (
+                            attention_mask.size() == input_ids.size()
+                        ):
+                            yield {
+                                "input_ids": input_ids,
+                                "labels": labels,
+                                "attention_mask": attention_mask,
+                                "position_ids": position_ids,
+                            }
+                        else:
+                            LOG.warning(
+                                "Dropping batch due to tensor size mismatch "
+                                f"input_ids: {input_ids.size()}, "
+                                f"labels: {labels.size()}, "
+                                f"attention_mask: {attention_mask.size()}"
+                            )
+                    buffer = {
+                        "input_ids": [],
+                        "attention_mask": [],
+                        "labels": [],
+                        "position_ids": [],
+                    }
+                    buffer_len = 0
+                    idx = 1
+
+                if example:
+                    # FIXME
+                    # just going to drop data points that are too long
+                    if len(example["input_ids"]) <= self.seq_length:
+                        input_ids = example["input_ids"]
+                        attention_mask = example["attention_mask"]
+                        labels = example["labels"]
+
+                        if add_concat_token:
+                            input_ids.append(self.concat_token_id)
+                            attention_mask.append(1)
+                            labels.append(self.concat_token_id)
+
+                        input_ids_with_concat = torch.tensor(
+                            input_ids, dtype=self.tokens_dtype
+                        )
+                        attention_mask_with_concat = torch.tensor(
+                            [idx * m for m in attention_mask], dtype=torch.int16
+                        )
+                        labels_with_concat = torch.tensor(
+                            labels, dtype=self.tokens_dtype
+                        )
+                        position_ids = torch.arange(
+                            len(input_ids), dtype=self.tokens_dtype
+                        )
+
+                        buffer["input_ids"].append(input_ids_with_concat)
+                        buffer["attention_mask"].append(attention_mask_with_concat)
+                        buffer["labels"].append(labels_with_concat)
+                        buffer["position_ids"].append(position_ids)
+                        buffer_len += len(input_ids)

src/axolotl/integrations/base.py

@@ -142,7 +142,7 @@ class BasePlugin:
             model: The loaded model.
         """
 
-    def get_trainer_cls(self, cfg: DictDefault) -> type[Trainer] | None:
+    def get_trainer_cls(self, cfg: DictDefault) -> Trainer | None:
         """Returns a custom class for the trainer.
 
         Args:

src/axolotl/integrations/config.py

@@ -20,8 +20,8 @@ from typing import Any, Dict, List, Type
 
 from axolotl.utils.schemas.config import (
     AxolotlConfigWCapabilities as AxolotlConfigWCapabilitiesBase,
-    AxolotlInputConfig as AxolotlInputConfigBase,
 )
+from axolotl.utils.schemas.config import AxolotlInputConfig as AxolotlInputConfigBase
 
 
 def merge_input_args():

src/axolotl/integrations/cut_cross_entropy/README.md

@@ -19,7 +19,7 @@ python scripts/cutcrossentropy_install.py | sh
 
 - If you are installing from pip
 ```bash
-pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c5aa3ef"
+pip3 uninstall -y cut-cross-entropy && pip3 install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5"
 ```
 
 ## Usage
@@ -34,7 +34,6 @@ plugins:
 - arcee
 - cohere
 - cohere2
-- deepseek_v3
 - gemma
 - gemma2
 - gemma3
@@ -43,7 +42,6 @@ plugins:
 - gemma3n_text
 - glm
 - glm4
-- glm4_moe
 - gpt_oss
 - granite
 - granitemoe
@@ -65,9 +63,7 @@ plugins:
 - qwen2_5_vl
 - qwen3
 - qwen3_moe
-- qwen3_next
 - smollm3
-- seed_oss
 - voxtral
 
 ## Citation

src/axolotl/integrations/cut_cross_entropy/__init__.py

@@ -35,7 +35,7 @@ LOG = get_logger(__name__)
 
 _CCE_INSTALL_MESSAGE = (
     "Please install Axolotl's fork of cut_cross_entropy with transformers support using "
-    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c5aa3ef"`'
+    '`pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@c6a32c5"`'
 )
 
 

src/axolotl/integrations/diffusion/README.md

@@ -1,154 +0,0 @@
-# Diffusion LM Training Plugin for Axolotl
-
-This plugin enables diffusion language model training using an approach inspired by
-LLaDA (Large Language Diffusion Models) within Axolotl.
-
-## Overview
-
-LLaDA is a diffusion-based approach to language model training that uses:
-- **Random token masking** during training instead of next-token prediction
-- **Bidirectional attention** to allow the model to attend to the full context
-- **Importance weighting** based on masking probabilities for stable training
-
-This approach can lead to more robust language models with better understanding of
-bidirectional context.
-
-## Installation
-
-The plugin is included with Axolotl. See our
-[installation docs](https://docs.axolotl.ai/docs/installation.html).
-
-## Quickstart
-
-Train with an example config (Llama‑3.2 1B):
-   - Pretrain: `axolotl train examples/llama-3/diffusion-3.2-1b-pretrain.yaml`
-   - SFT: `axolotl train examples/llama-3/diffusion-3.2-1b-sft.yaml`
-
-### Basic Configuration
-
-You can also modify your existing configs to enable / customize diffusion training.
-
-Add the following to your Axolotl config:
-
-```yaml
-# Enable diffusion LM training plugin
-plugins:
-  - axolotl.integrations.diffusion.DiffusionPlugin
-```
-
-And, configure the nested `diffusion` block (defaults shown):
-
-```yaml
-diffusion:
-  noise_schedule: linear  # or "cosine"
-  min_mask_ratio: 0.1
-  max_mask_ratio: 0.9
-  num_diffusion_steps: 128
-  eps: 1e-3
-  importance_weighting: true
-
-  # Mask token (training auto-adds if missing, avoid pad/eos)
-  mask_token_str: "<|diffusion_mask|>"
-  # Or use an existing special token id (e.g., 128002 for Llama-3.x)
-  # mask_token_id: 128002
-
-  # Sample generation during training (optional)
-  generate_samples: true
-  generation_interval: 100
-  num_generation_samples: 3
-  generation_steps: 128
-  generation_temperature: 0.0
-  generation_max_length: 100
-```
-
-## Supported Models
-
-Any models that support 4D attention masks should work out of the box. If not, please
-create an [issue](https://github.com/axolotl-ai-cloud/axolotl/issues) or open a
-[PR](https://github.com/axolotl-ai-cloud/axolotl/compare)!
-
-## How It Works
-
-### Random Masking
-During training, tokens are randomly masked:
-- Sample timestep `t` uniformly from [0, 1]
-- Calculate masking probability: `p = (1 - eps) * t + eps`
-- Randomly mask tokens with probability `p`
-
-### Diffusion Loss
-
-Loss is computed only on masked tokens with (optional) importance weighting:
-
-```python
-loss = sum(cross_entropy(pred, target) / p_mask) / total_tokens
-```
-
-## Sample Generation
-
-When `diffusion.generate_samples: true`, the plugin generates samples during training:
-
-```
-Sample 1:
-   Original (45 tokens): The quick brown fox jumps over the lazy dog...
-   Masked (18/45 tokens, 40.0%): The [MASK] [MASK] fox [MASK] over [MASK] lazy [MASK]...
-   Generated: The quick brown fox jumps over the lazy dog...
-```
-
-Samples are logged to console and wandb (if enabled).
-
-## Inference
-
-Diffusion inference is integrated into the standard Axolotl CLI. Use the same config
-you trained with and run:
-
-```
-axolotl inference path/to/your-config.yaml
-```
-
-Optionally, pass `--gradio` to use a simple web interface.
-
-Interactive controls (prefix the prompt with commands):
-- `:complete N` → completion mode with N new masked tokens appended (default 64)
-- `:mask R` → random masking mode with target mask ratio R in [0.0, 1.0]
-
-Example session:
-
-```
-================================================================================
-Commands:
-:complete N -> completion mode with N tokens (default 64)
-:mask R     -> random masking with ratio R (0.0–1.0)
-================================================================================
-Give me an instruction (Ctrl + D to submit):
-
-:mask 0.4 The quick brown fox jumps over the lazy dog
-
-Masked (40.0%):
-The [MASK] brown [MASK] jumps over the [MASK] dog
-
-Generated:
-The quick brown fox jumps over the loud dog
-```
-
-## Metrics and Monitoring
-
-The plugin adds (or modifies) several metrics to track diffusion training:
-
-- `train/loss`: Weighted diffusion loss
-- `train/accuracy`: Accuracy on masked tokens
-- `train/mask_ratio`: Average fraction of tokens masked
-- `train/num_masked_tokens`: Number of tokens masked
-- `train/avg_p_mask`: Average masking probability
-- `train/ce_loss`: Unweighted cross-entropy loss
-- `train/importance_weight_avg`: Average importance weight
-
-## Limitations
-
-- No flash attention support
-- No RL training support
-
-## References
-
-- [LLaDA Paper](https://arxiv.org/abs/2404.10406)
-- [Axolotl Documentation](https://docs.axolotl.ai/)
-- [API reference for plugin](https://docs.axolotl.ai/docs/api/integrations.diffusion.args.html#axolotl.integrations.diffusion.args)

src/axolotl/integrations/diffusion/__init__.py

@@ -1,19 +0,0 @@
-"""Diffusion LM training plugin init."""
-
-from .args import DiffusionArgs, DiffusionConfig
-from .callbacks import DiffusionGenerationCallback
-from .generation import generate
-from .plugin import DiffusionPlugin
-from .trainer import DiffusionTrainer
-from .utils import create_bidirectional_attention_mask, resolve_mask_token_id
-
-__all__ = [
-    "DiffusionArgs",
-    "DiffusionPlugin",
-    "DiffusionTrainer",
-    "generate",
-    "resolve_mask_token_id",
-    "create_bidirectional_attention_mask",
-    "DiffusionGenerationCallback",
-    "DiffusionConfig",
-]

src/axolotl/integrations/diffusion/args.py

@@ -1,95 +0,0 @@
-"""Config args for diffusion LM training (nested under `diffusion:`)."""
-
-from __future__ import annotations
-
-from typing import Literal
-
-from pydantic import BaseModel, Field, model_validator
-
-
-class DiffusionConfig(BaseModel):
-    """Nested diffusion configuration available under the `diffusion` key."""
-
-    # Noise schedule config
-    noise_schedule: Literal["linear", "cosine"] = Field(
-        default="linear", description="Type of noise schedule for diffusion training"
-    )
-    min_mask_ratio: float = Field(
-        default=0.1,
-        ge=0.0,
-        le=1.0,
-        description="Minimum masking ratio for diffusion noise schedule",
-    )
-    max_mask_ratio: float = Field(
-        default=0.9,
-        ge=0.0,
-        le=1.0,
-        description="Maximum masking ratio for diffusion noise schedule",
-    )
-    num_diffusion_steps: int = Field(
-        default=128, ge=1, description="Number of diffusion timesteps"
-    )
-    eps: float = Field(
-        default=1e-3,
-        ge=0.0,
-        le=1.0,
-        description="Epsilon value for minimum masking probability in forward process",
-    )
-
-    # Training config
-    importance_weighting: bool = Field(
-        default=True,
-        description="Apply importance weighting to loss based on masking probability",
-    )
-    mask_token_id: int | None = Field(
-        default=None,
-        description=(
-            "Token ID to use for masking. Unset by default; can use one of the "
-            "tokenizer's special tokens here."
-        ),
-    )
-    mask_token_str: str | None = Field(
-        default=None,
-        description=(
-            "Token string to use as a mask. If `mask_token_id` is invalid or unset, "
-            "this token will be ensured to exist as an additional special token and "
-            "used. If absent, a default '<|diffusion_mask|>' will be added."
-        ),
-    )
-
-    # Sample generation config
-    generate_samples: bool = Field(
-        default=True, description="Enable sample generation during training"
-    )
-    generation_interval: int = Field(
-        default=100, ge=1, description="Generate samples every N steps"
-    )
-    num_generation_samples: int = Field(
-        default=3, ge=1, description="Number of samples to generate each time"
-    )
-    generation_steps: int = Field(
-        default=128, ge=1, description="Number of diffusion steps for generation"
-    )
-    generation_temperature: float = Field(
-        default=0.0,
-        ge=0.0,
-        description="Temperature for generation sampling (0.0 = deterministic)",
-    )
-    generation_max_length: int = Field(
-        default=100, ge=1, description="Maximum sequence length for generation"
-    )
-
-    @model_validator(mode="after")
-    def _validate_mask_ratios(self) -> "DiffusionConfig":
-        if self.min_mask_ratio > self.max_mask_ratio:
-            raise ValueError("min_mask_ratio must be ≤ max_mask_ratio")
-        return self
-
-
-class DiffusionArgs(BaseModel):
-    """Plugin entry that exposes the nested `diffusion` block to the core config."""
-
-    diffusion: DiffusionConfig = Field(
-        default_factory=DiffusionConfig,
-        description="Diffusion training configuration. Only nested block is supported.",
-    )

src/axolotl/integrations/diffusion/callbacks.py

@@ -1,174 +0,0 @@
-"""Callbacks for diffusion training."""
-
-import logging
-import sys
-
-import wandb
-from colorama import Fore, Style
-from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
-from transformers.training_args import TrainingArguments
-
-from .generation import generate_samples
-
-# Simpler logger for more readable sample generation
-logger = logging.getLogger(__name__)
-if not logger.handlers:
-    handler = logging.StreamHandler(sys.stdout)
-    handler.setFormatter(logging.Formatter("%(message)s"))
-    logger.addHandler(handler)
-    logger.propagate = False
-logger.setLevel(logging.INFO)
-
-
-class DiffusionGenerationCallback(TrainerCallback):
-    """Callback for generating samples during diffusion training."""
-
-    def __init__(self, trainer):
-        self.trainer = trainer
-
-    def on_step_end(
-        self,
-        args: TrainingArguments,
-        state: TrainerState,
-        control: TrainerControl,
-        **kwargs,
-    ):
-        """Generate samples at specified intervals."""
-        if (
-            state.global_step > 0
-            and state.global_step % self.trainer.cfg.diffusion.generation_interval == 0
-        ):
-            if not self.trainer.state.is_world_process_zero:
-                return
-
-            # Use eval dataloader if available, otherwise use train dataloader
-            dataloader = None
-            try:
-                if getattr(self.trainer, "eval_dataset", None) is not None:
-                    dataloader = self.trainer.get_eval_dataloader()
-            except Exception:
-                dataloader = None
-            if dataloader is None:
-                dataloader = self.trainer.get_train_dataloader()
-
-            # Generate samples
-            diffusion_cfg = self.trainer.cfg.diffusion
-            samples = generate_samples(
-                model=self.trainer.model,
-                tokenizer=self.trainer.processing_class,
-                dataloader=dataloader,
-                num_generation_samples=diffusion_cfg.num_generation_samples,
-                max_length=diffusion_cfg.generation_max_length,
-                num_diffusion_steps=diffusion_cfg.generation_steps,
-                temperature=diffusion_cfg.generation_temperature,
-                mask_token_id=diffusion_cfg.mask_token_id,
-            )
-
-            # Log samples
-            self._log_samples(samples, state.global_step)
-
-    def _log_samples(self, samples: list, step: int):
-        """Log generated samples."""
-        if not samples:
-            return
-
-        logger.info("=" * 60)
-        logger.info("GENERATED SAMPLES")
-        logger.info("=" * 60)
-
-        for i, sample_data in enumerate(samples, 1):
-            original = sample_data["original"]
-            masked = sample_data["masked"]
-            generated = sample_data["generated"]
-            mask_ratio = sample_data["mask_ratio"]
-            masked_tokens = sample_data["masked_tokens"]
-            total_tokens = sample_data["total_tokens"]
-
-            logger.info(f"\nSample {i}:")
-            logger.info(f"\tOriginal ({total_tokens} tokens): {original}")
-            logger.info(
-                f"\tMasked ({masked_tokens}/{total_tokens} tokens, "
-                f"{mask_ratio:.1%}): {masked}"
-            )
-
-            try:
-                gen_ids = sample_data.get("generated_ids")
-                orig_ids = sample_data.get("orig_ids")
-                masked_positions = set(sample_data.get("masked_positions") or [])
-                if isinstance(gen_ids, list) and isinstance(orig_ids, list):
-                    styles: list[str] = []
-                    for i, tid in enumerate(gen_ids):
-                        if i in masked_positions:
-                            if i < len(orig_ids) and tid == orig_ids[i]:
-                                styles.append("green")
-                            elif i < len(orig_ids):
-                                styles.append("red")
-                            else:
-                                styles.append("normal")
-                        else:
-                            same = i < len(orig_ids) and tid == orig_ids[i]
-                            styles.append("dim" if same else "normal")
-
-                    spans: list[tuple[str, int, int]] = []
-                    if gen_ids:
-                        cur = styles[0]
-                        start = 0
-                        for i in range(1, len(gen_ids)):
-                            s = styles[i]
-                            if s != cur:
-                                spans.append((cur, start, i))
-                                cur, start = s, i
-                        spans.append((cur, start, len(gen_ids)))
-
-                    parts = []
-                    for style_name, a, b in spans:
-                        chunk_text = self.trainer.processing_class.decode(
-                            gen_ids[a:b], skip_special_tokens=False
-                        )
-                        if style_name == "green":
-                            parts.append(Fore.GREEN + chunk_text + Style.RESET_ALL)
-                        elif style_name == "red":
-                            parts.append(Fore.RED + chunk_text + Style.RESET_ALL)
-                        else:
-                            if style_name == "dim":
-                                parts.append(Style.DIM + chunk_text + Style.RESET_ALL)
-                            else:
-                                parts.append(chunk_text)
-                    logger.info("\tGenerated:\n%s", "".join(parts))
-                else:
-                    logger.info(f"\tGenerated: {generated}")
-            except Exception:
-                logger.info(f"\tGenerated: {generated}")
-
-        logger.info("=" * 60)
-
-        if self.trainer.cfg.use_wandb:
-            if wandb.run is not None:
-                wandb.log(
-                    {
-                        "generated_samples": wandb.Table(
-                            columns=[
-                                "step",
-                                "original",
-                                "masked",
-                                "generated",
-                                "mask_ratio",
-                                "masked_tokens",
-                                "total_tokens",
-                            ],
-                            data=[
-                                [
-                                    step,
-                                    sample["original"],
-                                    sample["masked"],
-                                    sample["generated"],
-                                    f"{sample['mask_ratio']:.1%}",
-                                    sample["masked_tokens"],
-                                    sample["total_tokens"],
-                                ]
-                                for sample in samples
-                            ],
-                        )
-                    },
-                    step=step,
-                )

src/axolotl/integrations/diffusion/generation.py

@@ -1,409 +0,0 @@
-"""Sample generation utilities for diffusion training."""
-
-import re
-from typing import Any, List, Literal, Optional
-
-import torch
-
-from axolotl.utils.logging import get_logger
-
-from .utils import create_bidirectional_attention_mask
-
-LOG = get_logger(__name__)
-
-
-def generate_samples(
-    model: torch.nn.Module,
-    tokenizer: Any,
-    dataloader: Optional[Any] = None,
-    num_generation_samples: int = 3,
-    max_length: int = 100,
-    num_diffusion_steps: int = 128,
-    temperature: float = 0.0,
-    mask_token_id: int = 32000,
-    mode: Literal["random", "completion"] = "random",
-    completion_tokens: int = 0,
-    target_mask_ratio: Optional[float] = None,
-) -> List[dict]:
-    """
-    Generate text samples using the diffusion model by randomly masking sequences from
-    the given dataset and running the reverse diffusion process.
-
-    Args:
-        model: The wrapped or unwrapped model
-        tokenizer: Tokenizer for encoding/decoding
-        dataloader: Validation dataloader (for sampling sequences)
-        num_generation_samples: Number of samples to generate
-        max_length: Maximum length of sequences to use
-        num_diffusion_steps: Number of diffusion steps for generation
-        temperature: Temperature for sampling (0.0 = deterministic)
-        mask_token_id: Token ID used for masking
-
-    Returns:
-        List of dictionaries with original text, masked text, and generated text
-    """
-    if dataloader is None:
-        LOG.warning("No validation dataloader provided, cannot generate samples")
-        return []
-
-    unwrapped_model = model.module if hasattr(model, "module") else model
-    training = unwrapped_model.training
-    unwrapped_model.eval()
-
-    # Resolve device robustly (some modules don't expose `.device`)
-    device = getattr(unwrapped_model, "device", None)
-    if device is None:
-        try:
-            device = next(unwrapped_model.parameters()).device
-        except StopIteration:
-            device = torch.device("cpu")
-    generations = []
-
-    # Sample sequences from validation dataset
-    sampled_sequences = _sample_sequences_from_dataloader(
-        dataloader, num_generation_samples, max_length, device
-    )
-    LOG.info(f"Sampled {len(sampled_sequences)} sequences from validation dataset")
-
-    # Generate samples using reverse diffusion process
-    with torch.no_grad():
-        for sample in sampled_sequences:
-            if isinstance(sample, dict):
-                original_sequence = sample.get("input_ids")
-                labels_seq = sample.get("labels")
-                attn_seq = sample.get("attention_mask")
-            else:
-                original_sequence = sample
-                labels_seq = None
-                attn_seq = None
-            generation_result = generate(
-                unwrapped_model,
-                tokenizer,
-                original_sequence,
-                num_diffusion_steps,
-                temperature,
-                mask_token_id,
-                mode=mode,
-                completion_tokens=completion_tokens,
-                target_mask_ratio=target_mask_ratio,
-                labels=labels_seq,
-                attention_mask=attn_seq,
-            )
-            generations.append(generation_result)
-
-    # Restore prior training state
-    if training:
-        unwrapped_model.train()
-    else:
-        unwrapped_model.eval()
-
-    return generations
-
-
-def _sample_sequences_from_dataloader(
-    dataloader: Any, num_samples: int, max_length: int, device: torch.device
-) -> List[Any]:
-    """Sample sequences from validation dataloader."""
-    sampled_sequences: list[dict[str, torch.Tensor] | torch.Tensor] = []
-    sample_count = 0
-
-    # Skip a random number of batches (we could be more clever about this)
-    skip_batches = torch.randint(0, 10, (1,)).item()
-    batch_count = 0
-
-    for batch in dataloader:
-        # Skip some batches for variety
-        if batch_count < skip_batches:
-            batch_count += 1
-            continue
-
-        if sample_count >= num_samples:
-            break
-
-        batch_count += 1
-        input_ids = batch["input_ids"]
-        attention_mask = batch.get("attention_mask")
-        labels = batch.get("labels")
-
-        # Randomly sample from sequences in this batch
-        batch_indices = torch.randperm(input_ids.size(0)).tolist()
-
-        for i in batch_indices:
-            if sample_count >= num_samples:
-                break
-
-            # Get actual sequence length (non-padded)
-            if attention_mask is not None:
-                seq_len = attention_mask[i].sum().item()
-            else:
-                seq_len = input_ids.size(1)
-
-            if seq_len < 10:
-                continue
-
-            # Determine truncation length
-            max_total = min(seq_len, max_length)
-            if labels is not None:
-                labels_i = labels[i][:seq_len]
-                answer_mask = labels_i != -100
-                if not answer_mask.any():
-                    # No answer tokens; skip for SFT masking
-                    continue
-                first_ans_idx = int(
-                    torch.nonzero(answer_mask, as_tuple=False)[0].item()
-                )
-                prompt_len = first_ans_idx
-                if prompt_len >= max_total:
-                    # Prompt alone reaches cap; cannot include any answer
-                    continue
-                remaining_answer = int(answer_mask[prompt_len:].sum().item())
-                allowed_answer = max_total - prompt_len
-                take_answer = min(remaining_answer, allowed_answer)
-                if take_answer <= 0:
-                    continue
-                actual_length = prompt_len + take_answer
-            else:
-                actual_length = max_total
-
-            # Extract the (possibly truncated) sequence
-            sequence = input_ids[i][:actual_length].unsqueeze(0).to(device)
-            attn_seq = (
-                attention_mask[i][:actual_length].unsqueeze(0).to(device)
-                if attention_mask is not None
-                else None
-            )
-            if labels is not None:
-                labels_seq = labels[i][:actual_length].unsqueeze(0).to(device)
-                sampled_sequences.append(
-                    {
-                        "input_ids": sequence,
-                        "labels": labels_seq,
-                        "attention_mask": attn_seq,
-                    }
-                )
-            else:
-                if attn_seq is not None:
-                    sampled_sequences.append(
-                        {"input_ids": sequence, "attention_mask": attn_seq}
-                    )
-                else:
-                    sampled_sequences.append(sequence)
-            sample_count += 1
-
-    return sampled_sequences
-
-
-def generate(
-    model: torch.nn.Module,
-    tokenizer: Any,
-    original_sequence: torch.Tensor,
-    num_diffusion_steps: int,
-    temperature: float,
-    mask_token_id: int,
-    *,
-    mode: Literal["random", "completion"] = "random",
-    completion_tokens: int = 0,
-    target_mask_ratio: Optional[float] = None,
-    labels: Optional[torch.Tensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,
-) -> dict:
-    """Generate a single sample using reverse diffusion."""
-    # Get original text for comparison
-    original_text = tokenizer.decode(
-        original_sequence[0].cpu(), skip_special_tokens=True
-    )
-
-    # Build masked sequence
-    if (
-        labels is not None
-        and labels.numel() > 0
-        and (labels == -100).any()
-        and (labels != -100).any()
-    ):
-        # SFT case: completely mask all answer tokens (labels != -100)
-        total_tokens = original_sequence.size(1)
-        masked_indices = (labels != -100).to(dtype=torch.bool)
-        masked_sequence = original_sequence.clone()
-        masked_sequence[masked_indices] = mask_token_id
-        masked_tokens = int(masked_indices.sum().item())
-        mask_ratio = masked_tokens / max(int(total_tokens), 1)
-    elif mode == "completion" and completion_tokens > 0:
-        # Append mask tokens to the right for completion
-        total_tokens = original_sequence.size(1) + int(completion_tokens)
-        masked_indices = torch.zeros(
-            1, total_tokens, dtype=torch.bool, device=original_sequence.device
-        )
-        masked_indices[0, -int(completion_tokens) :] = True
-
-        append = torch.full(
-            (1, int(completion_tokens)), mask_token_id, device=original_sequence.device
-        )
-        masked_sequence = torch.cat([original_sequence, append], dim=1)
-        masked_tokens = int(completion_tokens)
-        mask_ratio = masked_tokens / total_tokens
-    else:
-        # Apply random masking with optional fixed ratio
-        total_tokens = original_sequence.size(1)
-        if target_mask_ratio is None:
-            min_ratio, max_ratio = 0.1, 0.7
-            target_mask_ratio = (
-                torch.rand(1).item() * (max_ratio - min_ratio) + min_ratio
-            )
-        target_masked_tokens = max(1, int(total_tokens * float(target_mask_ratio)))
-
-        # Create random mask indices
-        mask_positions = torch.randperm(total_tokens)[:target_masked_tokens]
-        masked_indices = torch.zeros(
-            1, total_tokens, dtype=torch.bool, device=original_sequence.device
-        )
-        masked_indices[0, mask_positions] = True
-
-        # Create masked sequence
-        masked_sequence = original_sequence.clone()
-        masked_sequence[masked_indices] = mask_token_id
-
-        # Calculate actual mask ratio
-        masked_tokens = masked_indices.sum().item()
-        mask_ratio = masked_tokens / total_tokens
-
-    # Get masked text for comparison
-    masked_text = tokenizer.decode(masked_sequence[0].cpu(), skip_special_tokens=False)
-    masked_text = _clean_masked_text(masked_text, tokenizer, mask_token_id)
-
-    # Run reverse diffusion process
-    sequence = masked_sequence.clone()
-    attention_mask = create_bidirectional_attention_mask(
-        sequence, attention_mask, sample_packing=attention_mask is not None
-    )
-    for step in range(num_diffusion_steps):
-        sequence = _diffusion_step(
-            model,
-            sequence,
-            step,
-            num_diffusion_steps,
-            temperature,
-            mask_token_id,
-            attention_mask,
-        )
-    generated_text = tokenizer.decode(sequence[0].cpu(), skip_special_tokens=True)
-
-    # Collect diagnostic info
-    final_ids = sequence[0].detach().cpu().tolist()
-    orig_ids_for_render = original_sequence[0].detach().cpu().tolist()
-    if masked_indices is not None:
-        masked_positions = (
-            torch.where(masked_indices[0])[0].detach().cpu().tolist()
-            if masked_indices.ndim == 2
-            else []
-        )
-    else:
-        masked_positions = []
-
-    result = {
-        "original": original_text,
-        "masked": masked_text,
-        "generated": generated_text,
-        "mask_ratio": mask_ratio,
-        "masked_tokens": masked_tokens,
-        "total_tokens": total_tokens,
-        "generated_ids": final_ids,
-        "masked_positions": masked_positions,
-        "orig_ids": orig_ids_for_render,
-        "formatted": (
-            f"Original: '{original_text}' → Masked: '{masked_text}' "
-            f"({mask_ratio:.1%}) → Generated: '{generated_text}'"
-        ),
-    }
-
-    return result
-
-
-def _clean_masked_text(masked_text: str, tokenizer: Any, mask_token_id: int) -> str:
-    """Clean up masked text for display."""
-    mask_token_repr = tokenizer.decode([mask_token_id], skip_special_tokens=False)
-    cleaned = masked_text.replace(mask_token_repr, "[MASK]")
-
-    # Remove literal special token strings
-    if hasattr(tokenizer, "special_tokens_map"):
-        for token_value in tokenizer.special_tokens_map.values():
-            if token_value and isinstance(token_value, str):
-                cleaned = cleaned.replace(token_value, "")
-
-    # Normalize whitespace but preserve newlines
-    cleaned = cleaned.replace("\r\n", "\n").replace("\r", "\n")
-    cleaned = re.sub(r"[ \t]+", " ", cleaned)
-    cleaned = "\n".join(line.rstrip() for line in cleaned.split("\n")).strip()
-    return cleaned
-
-
-def _diffusion_step(
-    model: torch.nn.Module,
-    sequence: torch.Tensor,
-    step: int,
-    num_diffusion_steps: int,
-    temperature: float,
-    mask_token_id: int,
-    attention_mask: torch.Tensor | None = None,
-) -> torch.Tensor:
-    """Perform a single diffusion step with remasking."""
-    # Only process if there are masked tokens remaining
-    current_mask = sequence == mask_token_id
-    if not current_mask.any():
-        return sequence
-
-    # Create or use provided attention mask
-    if attention_mask is None:
-        batch_size, seq_len = sequence.shape
-        attention_mask = torch.ones(
-            batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=sequence.device
-        )
-
-    # Forward pass
-    outputs = model(input_ids=sequence, attention_mask=attention_mask)
-    logits = outputs.logits
-
-    # Only sample at currently masked positions
-    if current_mask.any():
-        masked_logits = logits[current_mask]
-
-        # Apply temperature scaling
-        if temperature > 0:
-            scaled_logits = masked_logits / temperature
-        else:
-            scaled_logits = masked_logits
-
-        # Suppress mask token in outputs
-        scaled_logits[:, mask_token_id] = -float("inf")
-
-        if temperature > 0:
-            # Add Gumbel noise for sampling
-            gumbel_noise = -torch.log(
-                -torch.log(torch.rand_like(scaled_logits, dtype=torch.float32))
-            )
-            gumbel_logits = scaled_logits + gumbel_noise
-            predicted_tokens = torch.argmax(gumbel_logits, dim=-1)
-        else:
-            predicted_tokens = torch.argmax(scaled_logits, dim=-1)
-
-        # Calculate probabilities for confidence scoring
-        probs = torch.softmax(scaled_logits, dim=-1)
-        predicted_token_probs = probs[range(len(predicted_tokens)), predicted_tokens]
-
-        # Determine how many tokens to unmask this step
-        remaining_masked = current_mask.sum().item()
-        if step == num_diffusion_steps - 1:
-            num_to_unmask = remaining_masked
-        else:
-            unmask_ratio = 1.0 / (num_diffusion_steps - step)
-            num_to_unmask = max(1, int(remaining_masked * unmask_ratio))
-
-        # Select highest confidence predictions to unmask
-        if num_to_unmask >= remaining_masked:
-            sequence[current_mask] = predicted_tokens
-        else:
-            _, top_indices = predicted_token_probs.topk(num_to_unmask)
-            mask_positions = torch.where(current_mask)[1]
-            positions_to_unmask = mask_positions[top_indices]
-            sequence[0, positions_to_unmask] = predicted_tokens[top_indices]
-
-    return sequence

src/axolotl/integrations/diffusion/plugin.py

@@ -1,41 +0,0 @@
-"""Diffusion LM training plugin for Axolotl."""
-
-from peft import PeftModel
-from transformers import PreTrainedModel
-
-from axolotl.integrations.base import BasePlugin
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.logging import get_logger
-
-from .trainer import DiffusionTrainer
-
-LOG = get_logger(__name__)
-
-
-class DiffusionPlugin(BasePlugin):
-    """
-    Plugin for diffusion language model training.
-
-    This plugin enables diffusion-based training using the LLaDA approach, which uses
-    random masking and bidirectional attention to train language models.
-    """
-
-    def __init__(self):
-        super().__init__()
-        self.cfg = None
-
-    def get_input_args(self) -> str:
-        """Returns the pydantic model for LLaDA plugin arguments."""
-        return "axolotl.integrations.diffusion.DiffusionArgs"
-
-    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
-        """Perform actions after model is loaded."""
-        self.cfg = cfg
-
-    def get_trainer_cls(self, cfg: DictDefault) -> type[DiffusionTrainer] | None:
-        """Return custom trainer class for diffusion training."""
-        return DiffusionTrainer
-
-    def post_trainer_create(self, cfg: DictDefault, trainer: DiffusionTrainer):
-        """Configure trainer after creation."""
-        trainer.set_config(cfg)