diffusion custom models approach

* data_parallel_size in in VllmserveCliArgs

* moved to 43

.github/CONTRIBUTING.md

@@ -62,7 +62,7 @@ We welcome ideas for improvements and new features. To suggest an enhancement, o
 You can skip certain CI checks by including specific keywords in your commit messages:
 
 - `[skip ci]` or `skip ci` - Skips all CI checks for that commit
-- `[skip-e2e]` or `skip-e2e` - Skips only end-to-end tests while running other CI checks
+- `[skip-e2e]` or `skip-e2e` - Skips only end-to-end tests while running other CI checks. You may also include this in the title of your PR to disable end-to-end tests for the entire PR.
 
 ## Style Guidelines
 

.github/workflows/tests.yml

@@ -188,13 +188,44 @@ jobs:
         run: |
           find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
+  gate-skip-e2e:
+    needs: [pre-commit, pytest, pytest-sdist]
+    runs-on: ubuntu-latest
+    outputs:
+      skip: ${{ steps.compute.outputs.skip }}
+    steps:
+      - uses: actions/github-script@v7
+        id: compute
+        with:
+          script: |
+            const token = /\[skip-e2e\]/i;
+            let msg = '';
+            if (context.eventName === 'push') {
+              msg = context.payload.head_commit?.message || '';
+            } else if (context.eventName === 'pull_request') {
+              const { owner, repo } = context.repo;
+              const prNumber = context.payload.pull_request.number;
+              const commits = await github.paginate(
+                github.rest.pulls.listCommits,
+                { owner, repo, pull_number: prNumber, per_page: 100 }
+              );
+              msg = commits.at(-1)?.commit?.message || '';
+            }
+            const title = context.payload.pull_request?.title || '';
+            const body  = context.payload.pull_request?.body  || '';
+            const skip = token.test(msg) || token.test(title) || token.test(body);
+            core.setOutput('skip', String(skip));
+
   docker-e2e-tests-1st:
     # Run this job first as a gate for running the remainder of the test matrix
-    if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' && !github.event.pull_request.draft }}
+    if: >
+      github.repository_owner == 'axolotl-ai-cloud' &&
+      (github.event_name != 'pull_request' || !github.event.pull_request.draft) &&
+      needs.gate-skip-e2e.outputs.skip != 'true'
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
-    needs: [pre-commit, pytest, pytest-sdist]
+    needs: [pre-commit, pytest, pytest-sdist, gate-skip-e2e]
 
     strategy:
       fail-fast: false
@@ -240,13 +271,16 @@ jobs:
           modal run cicd.e2e_tests
 
   docker-e2e-tests:
-    if: ${{ github.repository_owner == 'axolotl-ai-cloud' && !github.event.pull_request.draft }}
+    if: >
+      github.repository_owner == 'axolotl-ai-cloud' &&
+      (github.event_name != 'pull_request' || !github.event.pull_request.draft) &&
+      needs.gate-skip-e2e.outputs.skip != 'true'
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
     # Only run the remainder of the matrix if the first e2e check passed;
     # this is to save on wasted compute costs for known failures that get caught in the first run
-    needs: [pre-commit, pytest, docker-e2e-tests-1st]
+    needs: [pre-commit, pytest, gate-skip-e2e, docker-e2e-tests-1st]
 
     strategy:
       fail-fast: false

docker/Dockerfile-base

@@ -37,7 +37,7 @@ WORKDIR /workspace
 
 RUN python3 -m pip install --upgrade pip && pip3 install -U packaging==23.2 setuptools==75.8.0 wheel && \
     python3 -m pip install --no-cache-dir -U torch==${PYTORCH_VERSION}+cu${CUDA} torchvision --extra-index-url https://download.pytorch.org/whl/cu$CUDA && \
-    python3 -m pip install --no-cache-dir "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
+    CAUSAL_CONV1D_FORCE_CXX11_ABI=TRUE CAUSAL_CONV1D_FORCE_BUILD=TRUE python3 -m pip install --no-cache-dir causal_conv1d==1.5.2 && \
     python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" && \
     python3 -m pip cache purge
 

docs/multimodal.qmd

@@ -13,10 +13,13 @@ format:
 - [Pixtral](#sec-pixtral)
 - [Llava-1.5](#sec-llava-15)
 - [Mistral-Small-3.1](#sec-mistral-small-31)
+- [Voxtral](#sec-voxtral)
 - [Gemma-3](#sec-gemma-3)
 - [Gemma-3n](#sec-gemma-3n)
 - [Qwen2-VL](#sec-qwen2-vl)
 - [Qwen2.5-VL](#sec-qwen25-vl)
+- [SmolVLM2](#sec-smolvlm2)
+- [LFM2-VL](#sec-lfm2-vl)
 
 ## Usage
 
@@ -31,7 +34,7 @@ skip_prepare_dataset: true
 remove_unused_columns: false  # leave columns in place as they are needed to handle image embeddings during training
 sample_packing: false  # not yet supported with multimodal
 
-chat_template:  # see in next section
+chat_template:  # see in next section if specified
 
 # example dataset
 datasets:
@@ -97,6 +100,16 @@ base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
 chat_template: mistral_v7_tekken
 ```
 
+### Voxtral {#sec-voxtral}
+
+::: {.callout-tip}
+Please make sure to install audio lib via `pip3 install librosa==0.11.0 'mistral_common[audio]==1.8.3'`
+:::
+
+```yaml
+base_model: mistralai/Voxtral-Mini-3B-2507
+```
+
 ### Gemma-3 {#sec-gemma-3}
 
 ::: {.callout-tip}
@@ -143,6 +156,26 @@ base_model: Qwen/Qwen2.5-VL-7B-Instruct
 chat_template: qwen2_vl  # same as qwen2-vl
 ```
 
+### SmolVLM2 {#sec-smolvlm2}
+
+::: {.callout-tip}
+Please make sure to install `num2words` via `pip3 install num2words==0.5.14`
+:::
+
+```yaml
+base_model: HuggingFaceTB/SmolVLM2-500M-Video-Instruct
+```
+
+### LFM2-VL {#sec-lfm2-vl}
+
+::: {.callout-warning}
+Please uninstall `causal-conv1d` via `pip3 uninstall -y causal-conv1d`
+:::
+
+```yaml
+base_model: LiquidAI/LFM2-VL-450M
+```
+
 ## Dataset Format
 
 For multi-modal datasets, we adopt an extended `chat_template` format similar to OpenAI's Message format.
@@ -181,6 +214,20 @@ You may need to install `librosa` via `pip3 install librosa==0.11.0`.
 
 :::
 
+### Video
+
+::: {.callout-warning}
+
+This is not well tested at the moment. We welcome contributors!
+
+:::
+
+For video loading, you can use the following keys within `content` alongside `"type": "video"`:
+
+- `"path": "/path/to/video.mp4"`
+- `"url": "https://example.com/video.mp4"`
+- `"video": np.ndarray | list[PIL.Image.Image] | torch.Tensor` (or list of the aforementioned)
+
 ### Example
 
 Here is an example of a multi-modal dataset:

examples/LiquidAI/README.md

@@ -0,0 +1,58 @@
+# Finetune Liquid Foundation Models 2 (LFM2) with Axolotl
+
+[Liquid Foundation Models 2 (LFM2)](https://huggingface.co/collections/LiquidAI/lfm2-686d721927015b2ad73eaa38) are a family of small, open-weight models from [Liquid AI](https://www.liquid.ai/) focused on quality, speed, and memory efficiency. Liquid AI released text-only [LFM2](https://huggingface.co/collections/LiquidAI/lfm2-686d721927015b2ad73eaa38) and text+vision [LFM2-VL](https://huggingface.co/collections/LiquidAI/lfm2-vl-68963bbc84a610f7638d5ffa) models.
+
+LFM2 features a new hybrid Liquid architecture with multiplicative gates, short-range convolutions, and grouped query attention, enabling fast training and inference.
+
+This guide shows how to fine-tune both the LFM2 and LFM2-VL models with Axolotl.
+
+## Getting Started
+
+1.  Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+    ```bash
+    # Ensure you have a compatible version of Pytorch installed
+    pip3 install packaging setuptools wheel ninja
+    pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+    ```
+
+2.  Run one of the finetuning examples below.
+
+    **LFM2**
+    ```bash
+    # FFT SFT (1x48GB @ 25GiB)
+    axolotl train examples/LiquidAI/lfm2-350m-fft.yaml
+    ```
+
+    **LFM2-VL**
+    ```bash
+    # LoRA SFT (1x48GB @ 2.7GiB)
+    axolotl train examples/LiquidAI/lfm2-vl-lora.yaml
+    ```
+
+### TIPS
+
+- **Installation Error**: If you encounter `ImportError: ... undefined symbol ...` or `ModuleNotFoundError: No module named 'causal_conv1d_cuda'`, the `causal-conv1d` package may have been installed incorrectly. Try uninstalling it:
+  ```bash
+  pip uninstall -y causal-conv1d
+  ```
+
+- **Dataset Loading**: Read more on how to load your own dataset in our [documentation](https://docs.axolotl.ai/docs/dataset_loading.html).
+- **Dataset Formats**:
+  - For LFM2 models, the dataset format follows the OpenAI Messages format as seen [here](https://docs.axolotl.ai/docs/dataset-formats/conversation.html#chat_template).
+  - For LFM2-VL models, Axolotl follows the multi-content Messages format. See our [Multimodal docs](https://docs.axolotl.ai/docs/multimodal.html#dataset-format) for details.
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+
+## Related Resources
+
+- [LFM2 Blog](https://www.liquid.ai/blog/liquid-foundation-models-v2-our-second-series-of-generative-ai-models)
+- [LFM2-VL Blog](https://www.liquid.ai/blog/lfm2-vl-efficient-vision-language-models)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/LiquidAI/lfm2-350m-fft.yaml

@@ -2,7 +2,6 @@ base_model: LiquidAI/LFM2-350M
 
 chunked_cross_entropy: true
 
-chat_template: tokenizer_default
 eot_tokens:
   - "<|im_end|>"
 datasets:

examples/LiquidAI/lfm2-vl-lora.yaml

@@ -0,0 +1,58 @@
+base_model: LiquidAI/LFM2-VL-450M
+trust_remote_code: true
+model_type: AutoModelForImageTextToText
+processor_type: AutoProcessor
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.language_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+logging_steps: 1
+flash_attention: true
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

examples/gpt-oss/README.md

@@ -33,13 +33,44 @@ Note: Memory usage taken from `device_mem_reserved(gib)` from logs.
 
 ### Training 120B
 
-On 8xH100s
+On 8xH100s, make sure you have ~3TB of free disk space. With each checkpoint clocking in at ~720GB, along with the base
+model, and final model output, you may need at least 3TB of free disk space to keep at least 2 checkpoints.
 
 ```bash
 # FFT SFT with offloading (8x80GB @ ~49GiB/GPU)
 axolotl train examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml
 ```
 
+ERRATA: Transformers saves the model Architecture prefixed with `FSDP` which needs to be manually renamed in `config.json`.
+See https://github.com/huggingface/transformers/pull/40207 for the status of this issue.
+
+```bash
+sed -i 's/FSDPGptOssForCausalLM/GptOssForCausalLM/g' ./outputs/gpt-oss-out/config.json
+```
+
+When using SHARDED_STATE_DICT with FSDP, the final checkpoint should automatically merge the sharded weights to your
+configured `output_dir`. However, if that step fails due to a disk space error, you can take an additional step to
+merge the sharded weights.  This step will automatically determine the last checkpoint directory and merge the sharded
+weights to `{output_dir}/merged`.
+
+```bash
+axolotl merge-sharded-fsdp-weights examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml
+mv ./outputs/gpt-oss-out/merged/* ./outputs/gpt-oss-out/
+```
+
+
+### Inferencing your fine-tuned model
+
+GPT-OSS support in vLLM does not exist in a stable release yet. See https://x.com/MaziyarPanahi/status/1955741905515323425
+for more information about using a special vllm-openai docker image for inferencing with vLLM.
+
+SGLang has 0-day support in main, see https://github.com/sgl-project/sglang/issues/8833 for infomation on installing
+SGLang from source. Once you've installed SGLang, run the following command to launch a SGLang server:
+
+```bash
+python3 -m sglang.launch_server --model ./outputs/gpt-oss-out/ --served-model-name axolotl/gpt-oss-120b --host 0.0.0.0 --port 8888 --tp 8
+```
+
 ### Tool use
 
 GPT-OSS has a comprehensive tool understanding. Axolotl supports tool calling datasets for Supervised Fine-tuning.

examples/gpt-oss/gpt-oss-120b-fft-fsdp2-offload.yaml

@@ -20,6 +20,7 @@ datasets:
 dataset_prepared_path: last_run_prepared
 val_set_size: 0
 output_dir: ./outputs/gpt-oss-out/
+save_total_limit: 2  # the 120B model can use up to 720GB of disk space per checkpoint, so let's only keep the last 2
 
 sequence_len: 4096
 sample_packing: true

examples/lfm2/README.md

@@ -1,7 +0,0 @@
-# Liquid Foundation Models 2
-
-LFM2 support in transformers exists in the main branch, but is not yet included in the transformers release.
-
-```bash
-pip install --upgrade --no-deps --force-reinstall git+https://github.com/huggingface/transformers.git
-```

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

@@ -0,0 +1,57 @@
+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:
+  - diffusion.DiffusionPlugin
+noise_schedule: cosine
+min_mask_ratio: 0.15
+max_mask_ratio: 0.85
+eps: 5e-4
+importance_weighting: true
+mask_token_id: 128002
+generate_samples: true
+generation_interval: 10
+
+output_dir: ./outputs/model-out
+
+sequence_len: 512
+sample_packing: true
+
+gradient_accumulation_steps: 8
+micro_batch_size: 4
+max_steps: 10000
+
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 3e-4
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+sdp_attention: true
+
+warmup_steps: 1000
+
+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-3.2-1b-sft.yaml

@@ -0,0 +1,58 @@
+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:
+  - diffusion.DiffusionPlugin
+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
+
+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
+
+optimizer: adamw_8bit
+lr_scheduler: cosine
+learning_rate: 1e-5
+
+bf16: auto
+tf32: true
+
+gradient_checkpointing: true
+resume_from_checkpoint:
+logging_steps: 1
+sdp_attention: true
+
+warmup_steps: 1000
+
+save_strategy: steps
+eval_strategy: steps
+save_steps: 500
+eval_steps: 500
+
+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/smolvlm2/README.md

@@ -0,0 +1,49 @@
+# Finetune SmolVLM2 with Axolotl
+
+[SmolVLM2](https://huggingface.co/collections/HuggingFaceTB/smolvlm2-smallest-video-lm-ever-67ab6b5e84bf8aaa60cb17c7) are a family of lightweight, open-source multimodal models from HuggingFace designed to analyze and understand video, image, and text content.
+
+These models are built for efficiency, making them well-suited for on-device applications where computational resources are limited. Models are available in multiple sizes, including 2.2B, 500M, and 256M.
+
+This guide shows how to fine-tune SmolVLM2 models with Axolotl.
+
+## Getting Started
+
+1.  Install Axolotl following the [installation guide](https://docs.axolotl.ai/docs/installation.html).
+
+    Here is an example of how to install from pip:
+    ```bash
+    # Ensure you have a compatible version of Pytorch installed
+    pip3 install packaging setuptools wheel ninja
+    pip3 install --no-build-isolation 'axolotl[flash-attn]>=0.12.0'
+    ```
+
+2. Install an extra dependency:
+
+    ```bash
+    pip3 install num2words==0.5.14
+    ```
+
+3.  Run the finetuning example:
+
+    ```bash
+    # LoRA SFT (1x48GB @ 6.8GiB)
+    axolotl train examples/smolvlm2/smolvlm2-2B-lora.yaml
+    ```
+
+## TIPS
+
+- **Dataset Format**: For video finetuning, your dataset must be compatible with the multi-content Messages format. For more details, see our documentation on [Multimodal Formats](https://docs.axolotl.ai/docs/multimodal.html#dataset-format).
+- **Dataset Loading**: Read more on how to prepare and load your own datasets in our [documentation](https://docs.axolotl.ai/docs/dataset_loading.html).
+
+## Optimization Guides
+
+- [Multi-GPU Training](https://docs.axolotl.ai/docs/multi-gpu.html)
+- [LoRA Optimizations](https://docs.axolotl.ai/docs/lora_optims.html)
+- [Multi-Node Training](https://docs.axolotl.ai/docs/multi-node.html)
+
+## Related Resources
+
+- [SmolVLM2 Blog](https://huggingface.co/blog/smolvlm2)
+- [Axolotl Docs](https://docs.axolotl.ai)
+- [Axolotl GitHub](https://github.com/axolotl-ai-cloud/axolotl)
+- [Axolotl Discord](https://discord.gg/7m9sfhzaf3)

examples/smolvlm2/smolvlm2-2B-lora.yaml

@@ -0,0 +1,56 @@
+base_model: HuggingFaceTB/SmolVLM2-2.2B-Instruct
+trust_remote_code: true
+processor_type: AutoProcessor
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.text_model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+logging_steps: 1
+flash_attention: true
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+weight_decay: 0.0
+
+# save_first_step: true  # uncomment this to validate checkpoint saving works with your config

requirements.txt

@@ -1,7 +1,7 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 
 # START section of dependencies that don't install on Darwin/MacOS
-bitsandbytes==0.46.1
+bitsandbytes==0.47.0
 # triton 3.4.0 is not compatible with CCE
 triton>=3.0.0,<3.4.0
 mamba-ssm==1.2.0.post1
@@ -14,7 +14,7 @@ packaging==23.2
 
 huggingface_hub>=0.33.0
 peft==0.17.0
-transformers==4.55.0
+transformers==4.55.2
 tokenizers>=0.21.1
 accelerate==1.10.0
 datasets==4.0.0

src/axolotl/cli/args.py

@@ -40,6 +40,12 @@ class VllmServeCliArgs:
         default=None,
         metadata={"help": "Number of tensor parallel workers to use."},
     )
+    data_parallel_size: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "Number of data parallel workers to use for vLLM serving. This controls how many model replicas are used for parallel inference."
+        },
+    )
     host: Optional[str] = field(
         default=None,  # nosec B104
         metadata={"help": "Host address to run the server on."},

src/axolotl/cli/merge_sharded_fsdp_weights.py

@@ -10,6 +10,7 @@ import fire
 import torch
 import torch.distributed.checkpoint as dist_cp
 import torch.distributed.checkpoint.format_utils as dist_cp_format_utils
+from accelerate import PartialState
 from accelerate.utils import (
     SAFE_WEIGHTS_INDEX_NAME,
     SAFE_WEIGHTS_NAME,
@@ -23,6 +24,7 @@ from torch.distributed.checkpoint.format_utils import _EmptyStateDictLoadPlanner
 
 from axolotl.cli.config import load_cfg
 from axolotl.utils.logging import get_logger
+from axolotl.utils.train import determine_last_checkpoint
 
 LOG = get_logger(__name__)
 
@@ -143,7 +145,6 @@ def merge_fsdp_weights(
         ValueError: If torch version < 2.3.0, or if `checkpoint_dir` does not exist.
     """
     checkpoint_dir_ = Path(checkpoint_dir)
-    from accelerate.state import PartialState
 
     if not is_torch_version(">=", "2.3.0"):
         raise ValueError("`merge_fsdp_weights` requires PyTorch >= 2.3.0`")
@@ -180,7 +181,6 @@ def merge_fsdp_weights(
         if remove_checkpoint_dir:
             LOG.info(f"Removing old checkpoint directory {checkpoint_dir_}")
             shutil.rmtree(checkpoint_dir_)
-    state.wait_for_everyone()
 
 
 def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
@@ -195,11 +195,32 @@ def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
     parsed_cfg = load_cfg(config, **kwargs)
 
     fsdp_dir = Path(parsed_cfg.output_dir) / "pytorch_model_fsdp_0"
+    if not fsdp_dir.exists():
+        checkpoint_dir = determine_last_checkpoint(parsed_cfg, update=False)
+        if checkpoint_dir:
+            fsdp_dir = Path(checkpoint_dir) / "pytorch_model_fsdp_0"
+        if not fsdp_dir.exists():
+            raise ValueError(
+                f"Could not find FSDP checkpoint `pytorch_model_fsdp_0` in {checkpoint_dir}"
+            )
+
+    output_path = str(Path(parsed_cfg.output_dir) / "merged")
     merge_fsdp_weights(
         checkpoint_dir=str(fsdp_dir),
-        output_path=str(Path(parsed_cfg.output_dir) / "merged"),
+        output_path=output_path,
         safe_serialization=True,
     )
+    state = PartialState()
+    state.wait_for_everyone()
+    LOG.info(
+        f"FSDP SHARDED_STATE_DICT weights successfully merged to: {output_path}",
+        main_process_only=True,
+    )
+    LOG.info(
+        "Merged weights are only the safetensors and doesn't include the model configuration "
+        f"or tokenizer which may be found in {parsed_cfg.output_dir}.",
+        main_process_only=True,
+    )
 
 
 if __name__ == "__main__":

src/axolotl/core/builders/causal.py

@@ -10,6 +10,7 @@ import transformers
 from transformers import (
     DataCollatorWithFlattening,
     EarlyStoppingCallback,
+    Trainer,
 )
 from trl.trainer.utils import RewardDataCollatorWithPadding
 
@@ -385,10 +386,11 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 **data_collator_kwargs,
             )
         sig = inspect.signature(trainer_cls)
-        if "processing_class" in sig.parameters:
+        if "processing_class" in sig.parameters or issubclass(trainer_cls, Trainer):
             trainer_kwargs["processing_class"] = self.tokenizer
         elif "tokenizer" in sig.parameters:
             trainer_kwargs["tokenizer"] = self.tokenizer
+
         if (
             trainer_cls not in [AxolotlRewardTrainer, AxolotlPRMTrainer]
             and self.cfg.datasets is not None

src/axolotl/core/trainers/__init__.py

@@ -5,7 +5,6 @@
 
 from .base import AxolotlTrainer
 from .dpo.trainer import AxolotlDPOTrainer
-from .grpo.trainer import AxolotlGRPOSequenceParallelTrainer, AxolotlGRPOTrainer
 from .mamba import AxolotlMambaTrainer
 from .trl import (
     AxolotlCPOTrainer,

src/axolotl/core/trainers/base.py

@@ -82,7 +82,9 @@ class AxolotlTrainer(
         super().__init__(*_args, **kwargs)
 
         self.train_data_collator = self.data_collator
-        self._stored_metrics = defaultdict(lambda: defaultdict(list))
+        self._stored_metrics = defaultdict(
+            lambda: defaultdict(lambda: {"values": [], "reduction": "mean"})
+        )
         if self.args.orpo_alpha:
             self.loss_fct = torch.nn.CrossEntropyLoss(reduction="none")
 
@@ -573,9 +575,26 @@ class AxolotlTrainer(
         """
         # logs either has 'loss' or 'eval_loss'
         train_eval = "train" if "loss" in logs else "eval"
-        # Add averaged stored metrics to logs
-        for key, metrics in self._stored_metrics[train_eval].items():
-            logs[key] = torch.tensor(metrics).mean().item()
+
+        # Add reduced stored metrics to logs
+        for key, metric_data in self._stored_metrics[train_eval].items():
+            values = torch.tensor(metric_data["values"])
+            reduction_type = metric_data["reduction"]
+
+            if reduction_type == "mean":
+                logs[key] = values.mean().item()
+            elif reduction_type == "min":
+                logs[key] = values.min().item()
+            elif reduction_type == "max":
+                logs[key] = values.max().item()
+            elif reduction_type == "sum":
+                logs[key] = values.sum().item()
+            else:
+                raise NotImplementedError(
+                    "Metric reduction must be one of [mean, min, max, sum]"
+                )
+
+            logs[key] = round(logs[key], 4)
 
         if is_main_process():
             # Add memory usage
@@ -592,10 +611,27 @@ class AxolotlTrainer(
         return super().log(logs, start_time)
 
     def store_metrics(
-        self, metrics: dict[str, float], train_eval: Literal["train", "eval"] = "train"
+        self,
+        metrics: dict[str, float] | dict[str, tuple[int | float, str]],
+        train_eval: Literal["train", "eval"] = "train",
+        reduction: Literal["mean", "min", "max", "sum"] = "mean",
     ) -> None:
+        """
+        Store metrics with specified reduction type.
+
+        Args:
+            metrics: Dictionary of metric names to values, or metric names to (value,
+                reduction_type) tuples.
+            train_eval: Whether this is for training or evaluation.
+        """
         for key, value in metrics.items():
-            self._stored_metrics[train_eval][key].append(value)
+            if isinstance(value, tuple):
+                metric_value, metric_reduction = value
+            else:
+                metric_value, metric_reduction = value, reduction
+
+            self._stored_metrics[train_eval][key]["values"].append(metric_value)
+            self._stored_metrics[train_eval][key]["reduction"] = metric_reduction
 
     def _save_checkpoint(self, model, trial, **kwargs):
         # make sure the checkpoint dir exists, since trainer is flakey

src/axolotl/integrations/base.py

@@ -147,7 +147,7 @@ class BasePlugin:
         """
 
     # pylint: disable=unused-argument
-    def get_trainer_cls(self, cfg: DictDefault) -> Trainer | None:
+    def get_trainer_cls(self, cfg: DictDefault) -> type[Trainer] | None:
         """Returns a custom class for the trainer.
 
         Args:

src/axolotl/integrations/diffusion/README.md

@@ -0,0 +1,164 @@
+# Diffusion LM Training Plugin for Axolotl
+
+This plugin enables diffusion language model training using the LLaDA (Large Language
+And Diffusion Assistant) approach within the Axolotl framework.
+
+## 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 see 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. To use it, simply add the plugin configuration to
+your training config.
+
+## Quickstart
+
+### Basic Configuration
+
+Add the following to your Axolotl configuration YAML:
+
+```yaml
+# Enable diffusion LM training plugin
+plugins:
+  - axolotl.integrations.diffusion.DiffusionPlugin
+
+# Diffusion-specific configuration
+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_id: 128002
+
+# Sample generation (optional)
+generate_samples: true
+generation_interval: 100
+num_generation_samples: 3
+generation_steps: 128
+generation_temperature: 0.0
+generation_max_length: 100
+
+# Model configuration
+base_model: meta-llama/Llama-3.2-1B
+model_type: llama
+
+# Standard Axolotl configuration
+datasets:
+  - path: your_dataset
+    ...
+
+# Other config
+sequence_len: 1024
+micro_batch_size: 8
+gradient_accumulation_steps: 4
+learning_rate: 3e-4
+```
+
+## Supported Models
+
+Currently supported base model types:
+- **Llama** (meta-llama/Llama-*, etc.) - Uses `LlamaForDiffusionLM`
+- **Mistral** (mistralai/Mistral-*, etc.) - Uses `MistralForDiffusionLM`
+
+The plugin automatically creates custom model classes that inherit from the base model
+while adding diffusion training capabilities. This provides full compatibility with
+HuggingFace's ecosystem for saving, loading, and inference.
+
+## How It Works
+
+### Custom Model Architecture
+
+The plugin creates custom model classes (`LlamaForDiffusionLM`, `MistralForDiffusionLM`) that inherit from
+standard HuggingFace models. During training, these models:
+
+1. **Apply forward diffusion process**: Randomly mask tokens based on sampled timesteps
+2. **Use bidirectional attention**: Override causal attention with full bidirectional attention
+3. **Compute diffusion loss**: Calculate loss only on masked tokens with optional importance weighting
+
+### Random Masking
+During training, tokens are randomly masked based on a sampled timestep:
+- Sample timestep `t` uniformly from [0, 1]
+- Calculate masking probability: `p = (1 - eps) * t + eps`
+- Randomly mask tokens with probability `p`
+
+### Bidirectional Attention
+The models override causal attention with bidirectional attention:
+- Creates 4D attention masks allowing all-to-all attention
+- Maintains proper padding and sample packing masks
+- Compatible with standard HuggingFace attention implementations
+
+### Diffusion Loss
+
+Loss is computed only on masked tokens with (optional) importance weighting:
+
+```python
+loss = sum(cross_entropy(pred, target) / p_mask) / total_tokens
+```
+
+### Model Loading and Saving
+
+The custom models work seamlessly with HuggingFace's AutoModel system:
+
+```python
+from transformers import AutoModel, AutoConfig
+
+# Load a diffusion model
+model = AutoModel.from_pretrained("path/to/diffusion/model", trust_remote_code=True)
+
+# Save a diffusion model
+model.save_pretrained("path/to/save/diffusion/model")
+```
+
+During inference, the models behave like standard causal language models.
+
+## Sample Generation
+
+When `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).
+
+## Metrics and Monitoring
+
+The plugin adds 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
+
+## Benefits of Custom Model Approach
+
+✅ **Type Safety**: Full IDE support and type checking  
+✅ **HuggingFace Integration**: Works with AutoModel, Hub, pipelines  
+✅ **Maintainability**: Clean architecture, no monkey patching  
+✅ **Ecosystem Compatibility**: Standard save/load, PEFT support  
+✅ **Testing**: Easier to test and debug  
+
+## Limitations
+
+- **Model Support**: Currently limited to Llama and Mistral architectures
+- **Flash Attention**: Not yet optimized for flash attention
+- **Inference Speed**: Bidirectional attention is slower than causal for generation
+
+## References
+
+- [LLaDA Paper](https://arxiv.org/abs/2404.10406)
+- [Axolotl Documentation](https://docs.axolotl.ai/)

src/axolotl/integrations/diffusion/__init__.py

@@ -0,0 +1,26 @@
+"""Diffusion LM training plugin init."""
+
+from transformers import AutoConfig, AutoModel
+
+from .args import DiffusionArgs
+from .configuration import DiffusionConfig, LlamaForDiffusionConfig, MistralForDiffusionConfig
+from .models import LlamaForDiffusionLM, MistralForDiffusionLM
+from .plugin import DiffusionPlugin
+
+# Register custom configurations
+AutoConfig.register("llama_diffusion", LlamaForDiffusionConfig)
+AutoConfig.register("mistral_diffusion", MistralForDiffusionConfig)
+
+# Register custom models
+AutoModel.register(LlamaForDiffusionConfig, LlamaForDiffusionLM)
+AutoModel.register(MistralForDiffusionConfig, MistralForDiffusionLM)
+
+__all__ = [
+    "DiffusionArgs", 
+    "DiffusionPlugin",
+    "DiffusionConfig",
+    "LlamaForDiffusionConfig",
+    "MistralForDiffusionConfig", 
+    "LlamaForDiffusionLM",
+    "MistralForDiffusionLM",
+]

src/axolotl/integrations/diffusion/args.py

@@ -0,0 +1,70 @@
+"""Config args for diffusion LM training."""
+
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+
+class DiffusionArgs(BaseModel):
+    """Arguments for diffusion LM training plugin."""
+
+    # 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 = Field(
+        default=128002,
+        description=(
+            "Token ID to use for masking. Default is 128002 "
+            "(<|reserved_special_token_0|> for Llama 3.2)"
+        ),
+    )
+
+    # 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"
+    )

src/axolotl/integrations/diffusion/callbacks.py

@@ -0,0 +1,116 @@
+"""Callbacks for diffusion training."""
+
+import wandb
+from transformers.trainer_callback import TrainerCallback, TrainerControl, TrainerState
+from transformers.training_args import TrainingArguments
+
+from axolotl.utils.logging import get_logger
+
+from .generation import generate_samples
+
+LOG = get_logger(__name__)
+
+
+class DiffusionGenerationCallback(TrainerCallback):
+    """Callback for generating samples during diffusion training."""
+
+    def __init__(self, trainer):
+        self.trainer = trainer
+
+    # pylint: disable=unused-argument
+    def on_step_end(
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        """Generate samples at specified intervals."""
+        config = getattr(self.trainer, 'diffusion_config', self.trainer.args)
+        
+        if (
+            state.global_step > 0
+            and state.global_step % config.get('generation_interval', 100) == 0
+        ):
+            # Use eval dataloader if available, otherwise use train dataloader
+            if (
+                hasattr(self.trainer, "eval_dataset")
+                and self.trainer.eval_dataset is not None
+            ):
+                dataloader = self.trainer.get_eval_dataloader()
+            else:
+                dataloader = self.trainer.get_train_dataloader()
+
+            # Generate samples
+            samples = generate_samples(
+                model=self.trainer.model,
+                tokenizer=self.trainer.tokenizer,
+                dataloader=dataloader,
+                num_generation_samples=config.get('num_generation_samples', 3),
+                max_length=config.get('generation_max_length', 256),
+                num_diffusion_steps=config.get('generation_steps', 10),
+                temperature=config.get('generation_temperature', 1.0),
+                mask_token_id=config.get('mask_token_id', 32000),
+            )
+
+            # Log samples
+            self._log_samples(samples, state.global_step)
+
+    def _log_samples(self, samples: list, step: int):
+        """Log generated samples."""
+        if not samples:
+            return
+
+        LOG.info("=" * 60)
+        LOG.info("GENERATED SAMPLES")
+        LOG.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"]
+
+            LOG.info(f"\nSample {i}:")
+            LOG.info(f"\tOriginal ({total_tokens} tokens): {original}")
+            LOG.info(
+                f"\tMasked ({masked_tokens}/{total_tokens} tokens, "
+                f"{mask_ratio:.1%}): {masked}"
+            )
+            LOG.info(f"\tGenerated: {generated}")
+
+        LOG.info("=" * 60)
+
+        config = getattr(self.trainer, 'diffusion_config', self.trainer.args)
+        if config.get('use_wandb', False) and self.trainer.state.is_world_process_zero:
+            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/configuration.py

@@ -0,0 +1,71 @@
+"""Configuration classes for diffusion language models."""
+
+from transformers import LlamaConfig, MistralConfig
+
+
+class LlamaForDiffusionConfig(LlamaConfig):
+    """Configuration class for Llama models with diffusion training."""
+    
+    model_type = "llama_diffusion"
+    
+    def __init__(
+        self,
+        mask_token_id: int = 32000,
+        eps: float = 1e-3,
+        importance_weighting: bool = False,
+        sample_packing: bool = False,
+        min_mask_ratio: float = 0.0,
+        max_mask_ratio: float = 1.0,
+        noise_schedule: str = "linear",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        
+        # Diffusion-specific parameters
+        self.mask_token_id = mask_token_id
+        self.eps = eps
+        self.importance_weighting = importance_weighting
+        self.sample_packing = sample_packing
+        self.min_mask_ratio = min_mask_ratio
+        self.max_mask_ratio = max_mask_ratio
+        self.noise_schedule = noise_schedule
+
+
+class MistralForDiffusionConfig(MistralConfig):
+    """Configuration class for Mistral models with diffusion training."""
+    
+    model_type = "mistral_diffusion"
+    
+    def __init__(
+        self,
+        mask_token_id: int = 32000,
+        eps: float = 1e-3,
+        importance_weighting: bool = False,
+        sample_packing: bool = False,
+        min_mask_ratio: float = 0.0,
+        max_mask_ratio: float = 1.0,
+        noise_schedule: str = "linear",
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        
+        # Diffusion-specific parameters
+        self.mask_token_id = mask_token_id
+        self.eps = eps
+        self.importance_weighting = importance_weighting
+        self.sample_packing = sample_packing
+        self.min_mask_ratio = min_mask_ratio
+        self.max_mask_ratio = max_mask_ratio
+        self.noise_schedule = noise_schedule
+
+
+# Keep the base class for backward compatibility but mark as deprecated
+class DiffusionConfig(LlamaForDiffusionConfig):
+    """
+    Deprecated: Use LlamaForDiffusionConfig or MistralForDiffusionConfig instead.
+    """
+    
+    model_type = "diffusion"
+    
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)

src/axolotl/integrations/diffusion/generation.py

@@ -0,0 +1,269 @@
+"""Sample generation utilities for diffusion training."""
+
+import logging
+from typing import Any, List, Optional
+
+import torch
+
+logger = logging.getLogger(__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,
+) -> 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:
+        logger.warning("No validation dataloader provided, cannot generate samples")
+        return []
+
+    # Get the actual model (unwrap if needed)
+    unwrapped_model = model.module if hasattr(model, "module") else model
+    unwrapped_model.eval()
+    generations = []
+
+    # Sample sequences from validation dataset
+    sampled_sequences = _sample_sequences_from_dataloader(
+        dataloader, num_generation_samples, max_length, unwrapped_model.device
+    )
+    logger.info(f"Sampled {len(sampled_sequences)} sequences from validation dataset")
+
+    # Generate samples using reverse diffusion process
+    with torch.no_grad():
+        for original_sequence in sampled_sequences:
+            generation_result = _generate(
+                unwrapped_model,
+                tokenizer,
+                original_sequence,
+                num_diffusion_steps,
+                temperature,
+                mask_token_id,
+            )
+            generations.append(generation_result)
+
+    unwrapped_model.train()
+    return generations
+
+
+def _sample_sequences_from_dataloader(
+    dataloader: Any, num_samples: int, max_length: int, device: torch.device
+) -> List[torch.Tensor]:
+    """Sample sequences from validation dataloader."""
+    sampled_sequences = []
+    sample_count = 0
+
+    # Add randomness by skipping a random number of batches
+    skip_batches = torch.randint(0, 6, (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")
+
+        # 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)
+
+            # Limit sequence length to max_length
+            actual_length = min(seq_len, max_length)
+            if actual_length < 10:  # Skip very short sequences
+                continue
+
+            # Extract the sequence
+            sequence = input_ids[i][:actual_length].unsqueeze(0).to(device)
+            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,
+) -> 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
+    )
+
+    # Apply custom masking with random ratio (10% to 70%)
+    total_tokens = original_sequence.size(1)
+    min_ratio, max_ratio = 0.1, 0.7
+    target_mask_ratio = torch.rand(1).item() * (max_ratio - min_ratio) + min_ratio
+    target_masked_tokens = int(total_tokens * 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)
+    # Clean up mask token representation
+    masked_text = _clean_masked_text(masked_text, tokenizer, mask_token_id)
+
+    # Run reverse diffusion process
+    sequence = masked_sequence.clone()
+    for step in range(num_diffusion_steps):
+        sequence = _diffusion_step(
+            model, sequence, step, num_diffusion_steps, temperature, mask_token_id
+        )
+
+    # Get final generated text
+    generated_text = tokenizer.decode(sequence[0].cpu(), skip_special_tokens=True)
+
+    return {
+        "original": original_text,
+        "masked": masked_text,
+        "generated": generated_text,
+        "mask_ratio": mask_ratio,
+        "masked_tokens": masked_tokens,
+        "total_tokens": total_tokens,
+        "formatted": (
+            f"Original: '{original_text}' → Masked: '{masked_text}' "
+            f"({mask_ratio:.1%}) → Generated: '{generated_text}'"
+        ),
+    }
+
+
+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]")
+
+    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, "")
+
+    cleaned = " ".join(cleaned.split()).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,
+) -> 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 bidirectional attention mask for diffusion
+    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")
+
+        # Sample predictions
+        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:
+            # Deterministic sampling when temperature is 0
+            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/models.py

@@ -0,0 +1,426 @@
+"""Custom model classes for diffusion language models."""
+
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from transformers import LlamaForCausalLM, MistralForCausalLM
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration import LlamaForDiffusionConfig, MistralForDiffusionConfig
+
+
+class DiffusionModelMixin:
+    """Mixin class providing diffusion functionality to language models."""
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._special_token_ids = None
+        
+    def _cache_special_token_ids(self, tokenizer=None):
+        """Cache special token IDs to avoid repeated tokenizer access."""
+        if tokenizer is None:
+            self._special_token_ids = set()
+            return
+            
+        special_tokens = set()
+        
+        if hasattr(tokenizer, "bos_token_id") and tokenizer.bos_token_id is not None:
+            special_tokens.add(tokenizer.bos_token_id)
+        if hasattr(tokenizer, "eos_token_id") and tokenizer.eos_token_id is not None:
+            special_tokens.add(tokenizer.eos_token_id)
+        if hasattr(tokenizer, "pad_token_id") and tokenizer.pad_token_id is not None:
+            special_tokens.add(tokenizer.pad_token_id)
+            
+        self._special_token_ids = special_tokens
+    
+    def _forward_process(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        eps: float = 1e-3,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward noising process. A timestep is sampled along the process, and tokens are
+        masked with probability determined by the configured noise schedule.
+
+        Args:
+            input_ids: Input token ids [batch_size, seq_len].
+            attention_mask: Attention mask [batch_size, seq_len].
+            labels: Labels for SFT training [batch_size, seq_len].
+            eps: Small epsilon value for minimum masking probability.
+
+        Returns:
+            noisy_batch: Input with some tokens masked.
+            masked_indices: Boolean mask indicating which tokens were masked.
+            p_mask: Masking probabilities for each token [batch_size, seq_len].
+        """
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+
+        # Sample random timesteps for each sample in batch
+        t = torch.rand(batch_size, device=device)
+
+        # Calculate masking probability with epsilon
+        p_mask = (1 - eps) * t + eps  # [batch_size]
+        p_mask = p_mask[:, None].repeat(1, seq_len)  # [batch_size, seq_len]
+
+        # Don't mask padding tokens if attention_mask is provided
+        if attention_mask is not None:
+            valid_mask = attention_mask.bool()
+            p_mask = p_mask * valid_mask.float()
+
+        # Create mask to exclude special tokens
+        special_token_mask = torch.zeros_like(input_ids, dtype=torch.bool)
+        if self._special_token_ids:
+            for token_id in self._special_token_ids:
+                special_token_mask |= input_ids == token_id
+
+        # Create random mask based on p_mask
+        masked_indices = torch.rand((batch_size, seq_len), device=device) < p_mask
+        masked_indices = masked_indices & ~special_token_mask
+        if attention_mask is not None:
+            masked_indices = masked_indices & attention_mask.bool()
+
+        # For SFT data, only mask answer tokens
+        if labels is not None:
+            answer_mask = labels != -100
+            masked_indices = masked_indices & answer_mask
+
+        # Create masked input
+        mask_token_id = self.config.mask_token_id
+        noisy_batch = torch.where(masked_indices, mask_token_id, input_ids)
+
+        return noisy_batch, masked_indices, p_mask
+
+    def _create_bidirectional_attention_mask(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        """
+        Create bidirectional attention mask to override default causal masking. Handles
+        sample-packed sequences where different samples are identified by different
+        attention mask values.
+
+        Args:
+            input_ids: Input token ids [batch_size, seq_len].
+            attention_mask: Attention mask [batch_size, seq_len]
+
+        Returns:
+            bidirectional_mask: 4D attention mask [batch_size, 1, seq_len, seq_len].
+        """
+        batch_size, seq_len = input_ids.shape
+        device = input_ids.device
+
+        if attention_mask is None or not self.config.sample_packing:
+            return torch.ones(
+                batch_size, 1, seq_len, seq_len, dtype=torch.bool, device=device
+            )
+
+        # Create attention mask by comparing sample IDs element-wise
+        mask_i = attention_mask.unsqueeze(2)  # [batch_size, seq_len, 1]
+        mask_j = attention_mask.unsqueeze(1)  # [batch_size, 1, seq_len]
+
+        # Tokens can attend to each other if they have the same non-zero sample ID
+        bidirectional_mask = (mask_i == mask_j) & (mask_i > 0)
+
+        # Add head dimension: [batch_size, 1, seq_len, seq_len]
+        bidirectional_mask = bidirectional_mask.unsqueeze(1)
+
+        return bidirectional_mask
+
+    def _compute_diffusion_loss(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        logits: torch.Tensor | None = None,
+        masked_indices: torch.Tensor | None = None,
+        p_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        Compute diffusion loss given logits and masking information.
+
+        Args:
+            input_ids: Ground truth token ids [batch_size, seq_len].
+            attention_mask: Attention mask [batch_size, seq_len].
+            labels: Labels for SFT training [batch_size, seq_len].
+            logits: Model logits [batch_size, seq_len, vocab_size].
+            masked_indices: Boolean mask indicating which tokens were masked.
+            p_mask: Masking probabilities for each token [batch_size, seq_len].
+
+        Returns:
+            loss: Cross-entropy loss.
+        """
+        if masked_indices.sum() > 0:
+            valid_indices = torch.where(masked_indices)
+            batch_indices, seq_indices = valid_indices
+
+            masked_logits = logits[batch_indices, seq_indices]
+            masked_targets = input_ids[batch_indices, seq_indices]
+            masked_p_mask = p_mask[batch_indices, seq_indices]
+
+            # Compute cross-entropy loss without reduction
+            token_loss = F.cross_entropy(
+                masked_logits.float(), masked_targets, reduction="none"
+            )
+
+            if self.config.importance_weighting:
+                masked_p_mask = masked_p_mask.float()
+                weighted_loss = token_loss / masked_p_mask
+            else:
+                weighted_loss = token_loss
+
+            # Final loss: sum weighted losses, normalize
+            if labels is not None:
+                # For SFT data: normalize by answer length per sample
+                answer_mask = labels != -100
+                answer_lengths = answer_mask.sum(dim=1).float()  # [batch_size]
+
+                # Get batch indices for masked tokens
+                masked_batch_indices = batch_indices
+
+                # Sum losses per sample and divide by answer length
+                loss_per_sample = torch.zeros(
+                    input_ids.shape[0], device=input_ids.device
+                )
+                for i in range(input_ids.shape[0]):
+                    sample_mask = masked_batch_indices == i
+                    if sample_mask.sum() > 0:
+                        sample_loss = weighted_loss[sample_mask].sum()
+                        loss_per_sample[i] = sample_loss / answer_lengths[i]
+
+                loss = loss_per_sample.mean()
+            else:
+                # Original normalization for non-SFT data
+                loss = weighted_loss.sum() / (input_ids.shape[0] * input_ids.shape[1])
+        else:
+            loss = torch.tensor(0.0, device=input_ids.device, requires_grad=True)
+
+        return loss
+
+
+class LlamaForDiffusionLM(DiffusionModelMixin, LlamaForCausalLM):
+    """
+    Llama model for diffusion language modeling.
+    
+    This model extends LlamaForCausalLM with diffusion training capabilities,
+    including bidirectional attention and forward diffusion process.
+    """
+    
+    config_class = LlamaForDiffusionConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Initialize diffusion-specific attributes
+        self._special_token_ids = None
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+        
+    def set_tokenizer(self, tokenizer):
+        """Set tokenizer for special token handling."""
+        self._cache_special_token_ids(tokenizer)
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        """
+        Forward pass with diffusion training logic.
+        
+        During training, applies forward diffusion process and bidirectional attention.
+        During inference, behaves like standard causal language model.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if self.training and input_ids is not None:
+            # Apply diffusion process during training
+            original_input_ids = input_ids.clone()
+            
+            # Apply forward process to get noisy input
+            noisy_input_ids, masked_indices, p_mask = self._forward_process(
+                input_ids, attention_mask, labels, self.config.eps
+            )
+            
+            # Create bidirectional attention mask
+            bidirectional_attention_mask = self._create_bidirectional_attention_mask(
+                input_ids, attention_mask
+            )
+            
+            # Forward pass with noisy input and bidirectional attention
+            outputs = super().forward(
+                input_ids=noisy_input_ids,
+                attention_mask=bidirectional_attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=None,  # Don't use standard loss computation
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                cache_position=cache_position,
+                **kwargs,
+            )
+            
+            # Compute diffusion loss
+            loss = self._compute_diffusion_loss(
+                original_input_ids,
+                attention_mask,
+                labels,
+                outputs.logits,
+                masked_indices,
+                p_mask,
+            )
+            
+            if return_dict:
+                outputs.loss = loss
+                return outputs
+            else:
+                return (loss,) + outputs[1:]
+        else:
+            # Standard forward pass for inference
+            return super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                cache_position=cache_position,
+                **kwargs,
+            )
+
+
+class MistralForDiffusionLM(DiffusionModelMixin, MistralForCausalLM):
+    """
+    Mistral model for diffusion language modeling.
+    
+    This model extends MistralForCausalLM with diffusion training capabilities,
+    including bidirectional attention and forward diffusion process.
+    """
+    
+    config_class = MistralForDiffusionConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        
+        # Initialize diffusion-specific attributes
+        self._special_token_ids = None
+        
+        # Initialize weights and apply final processing
+        self.post_init()
+        
+    def set_tokenizer(self, tokenizer):
+        """Set tokenizer for special token handling."""
+        self._cache_special_token_ids(tokenizer)
+    
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[list[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        """
+        Forward pass with diffusion training logic.
+        
+        During training, applies forward diffusion process and bidirectional attention.
+        During inference, behaves like standard causal language model.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if self.training and input_ids is not None:
+            # Apply diffusion process during training
+            original_input_ids = input_ids.clone()
+            
+            # Apply forward process to get noisy input
+            noisy_input_ids, masked_indices, p_mask = self._forward_process(
+                input_ids, attention_mask, labels, self.config.eps
+            )
+            
+            # Create bidirectional attention mask
+            bidirectional_attention_mask = self._create_bidirectional_attention_mask(
+                input_ids, attention_mask
+            )
+            
+            # Forward pass with noisy input and bidirectional attention
+            outputs = super().forward(
+                input_ids=noisy_input_ids,
+                attention_mask=bidirectional_attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=None,  # Don't use standard loss computation
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                cache_position=cache_position,
+                **kwargs,
+            )
+            
+            # Compute diffusion loss
+            loss = self._compute_diffusion_loss(
+                original_input_ids,
+                attention_mask,
+                labels,
+                outputs.logits,
+                masked_indices,
+                p_mask,
+            )
+            
+            if return_dict:
+                outputs.loss = loss
+                return outputs
+            else:
+                return (loss,) + outputs[1:]
+        else:
+            # Standard forward pass for inference
+            return super().forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                cache_position=cache_position,
+                **kwargs,
+            )

src/axolotl/integrations/diffusion/plugin.py

@@ -0,0 +1,98 @@
+"""Diffusion LM training plugin for Axolotl."""
+
+from typing import TYPE_CHECKING
+
+from peft import PeftModel
+from transformers import AutoConfig, AutoModel, PreTrainedModel
+
+from axolotl.integrations.base import BasePlugin
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+from .callbacks import DiffusionGenerationCallback
+from .configuration import LlamaForDiffusionConfig, MistralForDiffusionConfig
+from .models import LlamaForDiffusionLM, MistralForDiffusionLM
+
+if TYPE_CHECKING:
+    from transformers import Trainer
+
+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 pre_model_load(self, cfg: DictDefault):
+        """Configure model loading to use diffusion model classes."""
+        # Map base model types to diffusion equivalents
+        base_model_type = cfg.get("model_type")
+        
+        if base_model_type == "llama":
+            # Create diffusion config from base config
+            diffusion_config = LlamaForDiffusionConfig(
+                mask_token_id=getattr(cfg, "mask_token_id", 32000),
+                eps=getattr(cfg, "eps", 1e-3),
+                importance_weighting=getattr(cfg, "importance_weighting", False),
+                sample_packing=getattr(cfg, "sample_packing", False),
+                min_mask_ratio=getattr(cfg, "min_mask_ratio", 0.0),
+                max_mask_ratio=getattr(cfg, "max_mask_ratio", 1.0),
+                noise_schedule=getattr(cfg, "noise_schedule", "linear"),
+            )
+            
+            # Override model type for loading
+            cfg.model_type = "llama_diffusion"
+            
+        elif base_model_type == "mistral":
+            # Create diffusion config from base config
+            diffusion_config = MistralForDiffusionConfig(
+                mask_token_id=getattr(cfg, "mask_token_id", 32000),
+                eps=getattr(cfg, "eps", 1e-3),
+                importance_weighting=getattr(cfg, "importance_weighting", False),
+                sample_packing=getattr(cfg, "sample_packing", False),
+                min_mask_ratio=getattr(cfg, "min_mask_ratio", 0.0),
+                max_mask_ratio=getattr(cfg, "max_mask_ratio", 1.0),
+                noise_schedule=getattr(cfg, "noise_schedule", "linear"),
+            )
+            
+            # Override model type for loading
+            cfg.model_type = "mistral_diffusion"
+        else:
+            LOG.warning(f"Diffusion plugin not implemented for model type: {base_model_type}")
+
+    def post_model_load(self, cfg: DictDefault, model: PreTrainedModel | PeftModel):
+        """Configure model after loading."""
+        self.cfg = cfg
+        
+        # Set tokenizer on diffusion models for special token handling
+        if hasattr(model, "set_tokenizer"):
+            # Get tokenizer from cfg if available
+            tokenizer = getattr(cfg, "tokenizer", None)
+            if tokenizer is not None:
+                model.set_tokenizer(tokenizer)
+
+    def add_callbacks_post_trainer(self, cfg: DictDefault, trainer: "Trainer"):
+        """Add diffusion-specific callbacks after trainer creation."""
+        callbacks = []
+        
+        # Store diffusion config on trainer for callbacks
+        trainer.diffusion_config = cfg
+        
+        # Add generation callback if enabled
+        if cfg.get("generate_samples", False):
+            generation_callback = DiffusionGenerationCallback(trainer)
+            callbacks.append(generation_callback)
+            
+        return callbacks

src/axolotl/loaders/constants.py

@@ -1,26 +1,13 @@
 """Shared constants for axolotl.loaders module"""
 
-from transformers import (
-    Gemma3ForConditionalGeneration,
-    Gemma3nForConditionalGeneration,
-    Llama4ForConditionalGeneration,
-    LlavaForConditionalGeneration,
-    Mistral3ForConditionalGeneration,
-    MllamaForConditionalGeneration,
-    Qwen2_5_VLForConditionalGeneration,
-    Qwen2VLForConditionalGeneration,
+from transformers import AutoModelForImageTextToText
+from transformers.models.auto.modeling_auto import (
+    MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES,
 )
 
-MULTIMODAL_AUTO_MODEL_MAPPING = {
-    "mllama": MllamaForConditionalGeneration,
-    "llama4": Llama4ForConditionalGeneration,
-    "llava": LlavaForConditionalGeneration,
-    "qwen2_vl": Qwen2VLForConditionalGeneration,
-    "qwen2_5_vl": Qwen2_5_VLForConditionalGeneration,
-    "mistral3": Mistral3ForConditionalGeneration,
-    "gemma3": Gemma3ForConditionalGeneration,
-    "gemma3n": Gemma3nForConditionalGeneration,
-}
+MULTIMODAL_AUTO_MODEL_MAPPING = dict(MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES)
+
+MULTIMODAL_AUTO_MODEL_MAPPING["lfm2-vl"] = AutoModelForImageTextToText
 
 try:
     from transformers import VoxtralForConditionalGeneration

src/axolotl/loaders/model.py

@@ -25,6 +25,7 @@ from peft import (
 from torch.distributed import DeviceMesh
 from transformers import (
     AutoModelForCausalLM,
+    AutoModelForImageTextToText,
     AutoModelForVision2Seq,
     AwqConfig,
     BitsAndBytesConfig,
@@ -212,6 +213,7 @@ class ModelLoader:
             self.model_kwargs["use_kernels"] = self.cfg.use_kernels
         self._set_quantization_config()
         self._set_attention_config()
+        self._check_model_requirements()
 
     def _apply_post_model_load_setup(self):
         """Configure the model after it has been loaded."""
@@ -432,6 +434,8 @@ class ModelLoader:
             self.auto_model_loader = MULTIMODAL_AUTO_MODEL_MAPPING.get(
                 self.model_config.model_type, AutoModelForVision2Seq
             )
+            if isinstance(self.auto_model_loader, str):
+                self.auto_model_loader = AutoModelForImageTextToText
 
     def _set_device_map_config(self):
         """Setup `device_map` according to config"""
@@ -628,6 +632,16 @@ class ModelLoader:
         if self.cfg.low_cpu_mem_usage:
             self.model_kwargs["low_cpu_mem_usage"] = True
 
+    def _check_model_requirements(self):
+        if self.cfg.model_config_type in ["lfm2-vl", "lfm2"]:
+            from transformers.utils.import_utils import is_causal_conv1d_available
+
+            if is_causal_conv1d_available():
+                raise ImportError(
+                    "The 'causal-conv1d' package is installed but causes compatibility issues with LFM2 models. "
+                    "Please uninstall it by running: `pip uninstall -y causal-conv1d`"
+                )
+
     def _configure_zero3_memory_efficient_loading(
         self,
     ) -> HfTrainerDeepSpeedConfig | None:
@@ -667,6 +681,23 @@ class ModelLoader:
 
         return hf_ds_cfg
 
+    def _load_model_from_config(self) -> PreTrainedModel:
+        """Load model with random initialization using from_config."""
+        if self.auto_model_loader in [AutoModelForCausalLM, AutoModelForVision2Seq]:
+            return self.auto_model_loader.from_config(config=self.model_config)
+        return self.auto_model_loader(config=self.model_config)
+
+    def _load_model_from_pretrained(self, model_loader_class=None) -> PreTrainedModel:
+        """Load model from pretrained weights."""
+        loader = model_loader_class or self.auto_model_loader
+        kwargs = {
+            **self.model_kwargs,
+            "config": self.model_config,
+            "trust_remote_code": self.cfg.trust_remote_code or False,
+            **self.model_kwargs,
+        }
+        return loader.from_pretrained(self.base_model, **kwargs)
+
     def _build_model(self) -> bool:
         """Load model, with load strategy depending on config."""
         skip_move_to_device = False
@@ -681,7 +712,8 @@ class ModelLoader:
         if self.is_fsdp_enabled:
             if self.cfg.fsdp_config.cpu_ram_efficient_loading:
                 skip_move_to_device = True
-                # Don't delete device_map for QLoRA + FSDP - it was set correctly in _set_device_map
+                # Don't delete device_map for QLoRA + FSDP - it was set correctly in
+                # _set_device_map
                 if (
                     "device_map" in self.model_kwargs
                     and not self.is_qlora_and_fsdp_enabled
@@ -710,6 +742,11 @@ class ModelLoader:
                 or self.cfg.qlora_sharded_model_loading
             )
         ):
+            if self.cfg.reinit_weights:
+                LOG.warning(
+                    "reinit_weights is not supported with sharded quantized loading. "
+                    "Loading from pretrained weights instead."
+                )
             quant_storage = self.cfg.torch_dtype
             quantization_config = getattr(
                 self.model_config, "quantization_config", None
@@ -725,33 +762,12 @@ class ModelLoader:
                 quantization_config=quantization_config,
             )
             skip_move_to_device = True
-        elif (
-            self.model_config.model_type in ["llama", "llama4"]
-            and not self.cfg.trust_remote_code
-            and not self.cfg.gptq
-        ):
-            # Please don't remove underscore binding without reading the fn docstring.
-            _ = self._configure_zero3_memory_efficient_loading()
-
-            # Load model with random initialization if specified
-            if self.cfg.random_init_weights:
-                # AutoModel classes support the from_config method
-                if self.auto_model_loader in [
-                    AutoModelForCausalLM,
-                    AutoModelForVision2Seq,
-                ]:
-                    self.model = self.auto_model_loader.from_config(
-                        config=self.model_config,
-                    )
-                else:
-                    self.model = self.auto_model_loader(config=self.model_config)
-            else:
-                self.model = self.auto_model_loader.from_pretrained(
-                    self.base_model,
-                    config=self.model_config,
-                    **self.model_kwargs,
-                )
         elif self.model_type == "MambaLMHeadModel":
+            if self.cfg.reinit_weights:
+                LOG.warning(
+                    "reinit_weights is not supported with MambaLMHeadModel. "
+                    "Loading from pretrained weights instead."
+                )
             # FIXME this is janky at best and hacked together to make it work
             MambaLMHeadModel = fix_mamba_attn_for_loss()  # pylint: disable=invalid-name
 
@@ -764,41 +780,27 @@ class ModelLoader:
                 self.base_model,
                 **self.model_kwargs,
             )
-        elif (
-            self.model_type
-            and self.model_type != "AutoModelForCausalLM"
-            and not self.cfg.trust_remote_code
-        ):
-            if self.cfg.gptq:
-                self.model = self.auto_model_loader.from_pretrained(
-                    self.base_model,
-                    config=self.model_config,
-                    trust_remote_code=self.cfg.trust_remote_code or False,
-                    **self.model_kwargs,
-                )
-            else:
-                self.model = getattr(transformers, self.model_type).from_pretrained(
-                    self.base_model,
-                    config=self.model_config,
-                    trust_remote_code=self.cfg.trust_remote_code or False,
-                    **self.model_kwargs,
-                )
-        elif self.cfg.gptq:
-            self.model = self.auto_model_loader.from_pretrained(
-                self.base_model,
-                config=self.model_config,
-                trust_remote_code=self.cfg.trust_remote_code or False,
-                **self.model_kwargs,
-            )
         else:
-            # Please don't remove underscore binding without reading the fn docstring.
+            # Please don't remove underscore binding without reading the fn docstring
             _ = self._configure_zero3_memory_efficient_loading()
-            self.model = self.auto_model_loader.from_pretrained(
-                self.base_model,
-                config=self.model_config,
-                trust_remote_code=self.cfg.trust_remote_code or False,
-                **self.model_kwargs,
-            )
+
+            if (
+                self.model_type
+                and self.model_type != "AutoModelForCausalLM"
+                and not self.cfg.trust_remote_code
+                and not self.cfg.gptq
+            ):
+                # Use model type from transformers
+                model_loader_class = getattr(transformers, self.model_type)
+            else:
+                # Use auto model loader (handles gptq and default cases)
+                model_loader_class = self.auto_model_loader
+
+            if self.cfg.reinit_weights:
+                self.model = self._load_model_from_config()
+            else:
+                self.model = self._load_model_from_pretrained(model_loader_class)
+
         if is_deepspeed_zero3_enabled():
             skip_move_to_device = True
 

src/axolotl/loaders/patch_manager.py

@@ -285,12 +285,10 @@ class PatchManager:
             and self.cfg.adapter == "qlora"
         ):
             from axolotl.monkeypatch.fsdp2_qlora import (
-                apply_bnb_torch_function_patch,
                 apply_init_sharded_param_patch,
                 apply_init_unsharded_param_patch,
             )
 
-            apply_bnb_torch_function_patch()
             apply_init_sharded_param_patch()
             apply_init_unsharded_param_patch()
 

src/axolotl/monkeypatch/fsdp2_qlora.py

@@ -9,73 +9,12 @@ Params4bit parameters.
 import importlib
 import inspect
 
-import torch
-from torch.nn import Parameter
-
 from axolotl.monkeypatch.utils import detab_code
 from axolotl.utils.logging import get_logger
 
 LOG = get_logger(__name__)
 
 
-def patched_torch_function(cls, func, types, args=(), kwargs=None):
-    """
-    Patched version of Params4bit.__torch_function__ for preserving Params4bit
-    class identity and attributes.
-    """
-    if kwargs is None:
-        kwargs = {}
-
-    if func in [torch.chunk, torch.split]:
-        tensor = args[0]
-        result = Parameter.__torch_function__(func, types, args, kwargs)
-
-        if isinstance(result, tuple):
-            return tuple(
-                cls(
-                    data=chunk,
-                    requires_grad=tensor.requires_grad,
-                    quant_state=tensor.quant_state,
-                    blocksize=tensor.blocksize,
-                    compress_statistics=tensor.compress_statistics,
-                    quant_type=tensor.quant_type,
-                    quant_storage=tensor.quant_storage,
-                    module=tensor.module,
-                    bnb_quantized=tensor.bnb_quantized,
-                )
-                for chunk in result
-            )
-
-        return cls(
-            data=result,
-            requires_grad=tensor.requires_grad,
-            quant_state=tensor.quant_state,
-            blocksize=tensor.blocksize,
-            compress_statistics=tensor.compress_statistics,
-            quant_type=tensor.quant_type,
-            quant_storage=tensor.quant_storage,
-            module=tensor.module,
-            bnb_quantized=tensor.bnb_quantized,
-        )
-
-    return Parameter.__torch_function__(func, types, args, kwargs)
-
-
-# pylint: disable=protected-access
-def apply_bnb_torch_function_patch():
-    """
-    Patch Params4bit.__torch_function__ using Axolotl-style approach.
-
-    Returns:
-        True if patching succeeded, False otherwise.
-    """
-    from bitsandbytes.nn.modules import Params4bit
-
-    Params4bit.__torch_function__ = classmethod(patched_torch_function)
-
-    LOG.info("Successfully patched Params4bit.__torch_function__")
-
-
 # pylint: disable=protected-access
 def apply_init_sharded_param_patch():
     """Apply patch to FSDPParam._init_sharded_param to support Params4bit."""

src/axolotl/monkeypatch/ring_attn/adapters/batch.py

@@ -20,12 +20,15 @@ from ring_flash_attn import ring_flash_attn_func
 from ring_flash_attn.adapters.hf_adapter import check_params
 from transformers.modeling_flash_attention_utils import is_flash_attn_greater_or_equal
 
-try:
+try:  # pylint: disable=duplicate-code
     from transformers.modeling_flash_attention_utils import _flash_supports_window
 except ImportError:
-    from transformers.modeling_flash_attention_utils import (
-        _flash_supports_window_size as _flash_supports_window,
-    )
+    try:
+        from transformers.modeling_flash_attention_utils import (
+            _flash_supports_window_size as _flash_supports_window,
+        )
+    except ImportError:
+        _flash_supports_window = True
 
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 

src/axolotl/monkeypatch/ring_attn/patch.py

@@ -15,12 +15,15 @@ import torch
 import torch.distributed as dist
 from torch.distributed import DeviceMesh
 
-try:
+try:  # pylint: disable=duplicate-code
     from transformers.modeling_flash_attention_utils import _flash_supports_window
 except ImportError:
-    from transformers.modeling_flash_attention_utils import (
-        _flash_supports_window_size as _flash_supports_window,
-    )
+    try:
+        from transformers.modeling_flash_attention_utils import (
+            _flash_supports_window_size as _flash_supports_window,
+        )
+    except ImportError:
+        _flash_supports_window = True
 
 from axolotl.monkeypatch.utils import get_cu_seqlens_from_pos_ids
 from axolotl.utils.logging import get_logger

src/axolotl/processing_strategies.py

@@ -6,7 +6,7 @@ from typing import Optional
 from PIL import Image, ImageOps
 from PIL.Image import Resampling
 from torch import Tensor, zeros_like
-from transformers import ProcessorMixin, VoxtralProcessor
+from transformers import ProcessorMixin, SmolVLMProcessor, VoxtralProcessor
 from transformers.image_utils import load_image
 
 from axolotl.utils.dict import remove_none_values
@@ -138,7 +138,7 @@ class ProcessingStrategy:
                     image_key = key
                     break
 
-            # if the image key exists, add the image to the first message
+            # if the image key exists, add the image to the first user message
             if image_key is not None and processed_example[image_key] is not None:
                 # TODO: check if it's normal to be single image only for common datasets
                 # From observation, it's usually a list of single image but some datasets may have several columns for images
@@ -179,26 +179,34 @@ class ProcessingStrategy:
 
                 # Look for any image type in the first message
                 # some dataset have an {type: "image"} in the first message
+                msg_ind_to_add = None
                 ind_to_add = None
+                first_user_idx = None
 
-                for i, content in enumerate(
-                    processed_example["messages"][0]["content"]
-                ):
-                    # Usually datasets created with image columns, don't have it in the messages itself
-                    if content["type"] == "image" and all(
-                        k not in content for k in ["image", "url", "path", "base64"]
+                for msg_idx, msg_content in enumerate(processed_example["messages"]):
+                    if first_user_idx is None and msg_content["role"] == "user":
+                        first_user_idx = msg_idx
+                    for i, content in enumerate(
+                        processed_example["messages"][msg_idx]["content"]
                     ):
-                        ind_to_add = i
-                        break
+                        # Usually datasets created with image columns, don't have it in the messages itself
+                        if content["type"] == "image" and all(
+                            k not in content for k in ["image", "url", "path", "base64"]
+                        ):
+                            msg_ind_to_add = msg_idx
+                            ind_to_add = i
+                            break
 
                 # If an image type is found, add the image to that index
-                if ind_to_add is not None:
-                    processed_example["messages"][0]["content"][ind_to_add][
-                        "image"
-                    ] = image_value
+                if ind_to_add is not None and msg_ind_to_add is not None:
+                    processed_example["messages"][msg_ind_to_add]["content"][
+                        ind_to_add
+                    ]["image"] = image_value
                 else:
-                    # if no image type is found, add it to end of the first message
-                    processed_example["messages"][0]["content"].append(
+                    # if no image type is found, add it to end of the first user message
+                    if first_user_idx is None:
+                        first_user_idx = 0
+                    processed_example["messages"][first_user_idx]["content"].append(
                         {
                             "type": "image",
                             "image": image_value,
@@ -395,6 +403,24 @@ class VoxtralProcessingStrategy(ProcessingStrategy):
         return labels
 
 
+class SmolVLM2ProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for SmolVLM2"""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+    ):
+        super().__init__(processor, chat_template, image_size, image_resize_algorithm)
+        self.image_token = "<image>"  # nosec
+
+        self.image_token_id = processor.tokenizer.additional_special_tokens_ids[
+            processor.tokenizer.additional_special_tokens.index(self.image_token)
+        ]
+
+
 def get_processing_strategy(
     processor: ProcessorMixin,
     chat_template,
@@ -402,32 +428,43 @@ def get_processing_strategy(
     image_size: int | tuple[int, int] | None = None,
     image_resize_algorithm: Resampling | None = None,
 ):
+    processing_kwargs = {
+        "processor": processor,
+        "chat_template": chat_template,
+        "image_size": image_size,
+        "image_resize_algorithm": image_resize_algorithm,
+    }
+
+    if chat_template_type in [None, "tokenizer_default"] and hasattr(
+        processor.tokenizer, "chat_template"
+    ):
+        processing_kwargs["chat_template"] = processor.tokenizer.chat_template
+
     if chat_template_type == "qwen2_vl":
         return Qwen2VLProcessingStrategy(
-            processor, chat_template, image_size, image_resize_algorithm
+            **processing_kwargs,
         )
     if chat_template_type == "gemma3":
         return Gemma3ProcessingStrategy(
-            processor, chat_template, image_size, image_resize_algorithm
+            **processing_kwargs,
         )
     if chat_template_type == "gemma3n":
         return Gemma3nProcessingStrategy(
-            processor, chat_template, image_size, image_resize_algorithm
-        )
-    if chat_template_type in [
-        "llama3_2_vision",
-        "llama4",
-        "llava",
-        "mistral_v7_tekken",
-        "pixtral",
-    ]:
-        return ProcessingStrategy(
-            processor, chat_template, image_size, image_resize_algorithm
+            **processing_kwargs,
         )
 
     if isinstance(processor, VoxtralProcessor):
         return VoxtralProcessingStrategy(
-            processor, chat_template, image_size, image_resize_algorithm
+            **processing_kwargs,
         )
 
-    raise ValueError(f"Unsupported chat template type: {chat_template_type}")
+    if isinstance(processor, SmolVLMProcessor):
+        return SmolVLM2ProcessingStrategy(
+            **processing_kwargs,
+        )
+
+    # llama3_2_vision, llama4, llava
+    # mistral_v7_tekken, pixtral, lfm2vl
+    return ProcessingStrategy(
+        **processing_kwargs,
+    )

src/axolotl/prompt_strategies/chat_template.py

@@ -129,13 +129,21 @@ class ChatTemplatePrompter(Prompter):
                 images=images,
                 return_tensors="pt",
             )
+            if hasattr(batch, "to_dict"):
+                batch = batch.to_dict()
+            else:
+                batch = dict(batch)
+
             # workaround since processor works in batches instead of single examples
+            out = {}
             for k, val in batch.items():
-                if k in ["pixel_values"]:
-                    batch[k] = val.tolist()
+                if hasattr(val, "tolist"):
+                    out[k] = (
+                        val.tolist() if k == "pixel_values" else val.squeeze(0).tolist()
+                    )
                 else:
-                    batch[k] = val.squeeze().tolist()
-            return batch
+                    out[k] = val
+            return out
 
         return self.tokenizer.apply_chat_template(
             conversation,
@@ -433,10 +441,13 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
                 tokenized_prompt["attention_mask"] = [1] * len(input_ids)
             else:
                 input_ids = tokenized_res["input_ids"]
-                tokenized_prompt = tokenized_res
+                tokenized_prompt = dict(tokenized_res)
 
             if not self.train_on_inputs:
-                user_prompt_len = len(prompt_ids)
+                if isinstance(prompt_ids, dict):
+                    user_prompt_len = len(prompt_ids["input_ids"])
+                else:
+                    user_prompt_len = len(prompt_ids)
                 labels = [-100] * user_prompt_len + input_ids[user_prompt_len:]
             else:
                 labels = input_ids

src/axolotl/prompt_tokenizers.py

@@ -75,7 +75,7 @@ class PromptTokenizingStrategy(abc.ABC):
     ) -> BatchEncoding:
         empty = BatchEncoding(data={"input_ids": [], "attention_mask": []})
         if not prompt:
-            LOG.warning("Empty text requested for tokenization.")
+            LOG.warning_once("Empty text requested for tokenization.")
             return empty
 
         result = self.tokenizer(

src/axolotl/train.py

@@ -4,11 +4,14 @@ from __future__ import annotations
 
 import importlib
 import inspect
+import json
 import os
+import shutil
 import signal
 import sys
 import typing
 import weakref
+from collections import OrderedDict
 from contextlib import ExitStack
 from pathlib import Path
 from typing import Any, Dict
@@ -38,6 +41,7 @@ from axolotl.utils.distributed import cleanup_distributed
 from axolotl.utils.freeze import freeze_layers_except
 from axolotl.utils.logging import get_logger
 from axolotl.utils.schemas.enums import RLType
+from axolotl.utils.train import determine_last_checkpoint
 from axolotl.utils.trainer import setup_trainer
 
 try:
@@ -46,7 +50,7 @@ except ImportError:
     BetterTransformer = None
 
 if typing.TYPE_CHECKING:
-    from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFRLTrainerBuilder
+    from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
 
 LOG = get_logger(__name__)
 
@@ -124,32 +128,6 @@ def setup_reference_model(
     return model_ref
 
 
-def determine_resume_checkpoint(cfg: DictDefault) -> str | None:
-    """
-    Determine the checkpoint to resume from based on configuration.
-
-    Args:
-        cfg: Dictionary mapping `axolotl` config keys to values.
-
-    Returns:
-        Path to the checkpoint to resume from, or `None` if not resuming.
-    """
-    if cfg.resume_from_checkpoint is None and cfg.auto_resume_from_checkpoints:
-        possible_checkpoints = [
-            str(cp) for cp in Path(cfg.output_dir).glob("checkpoint-*")
-        ]
-        if len(possible_checkpoints) > 0:
-            sorted_paths = sorted(
-                possible_checkpoints,
-                key=lambda path: int(path.split("-")[-1]),
-            )
-            cfg.resume_from_checkpoint = sorted_paths[-1]
-            LOG.info(
-                f"Using Auto-resume functionality to start with checkpoint at {cfg.resume_from_checkpoint}"
-            )
-    return cfg.resume_from_checkpoint
-
-
 def setup_signal_handler(
     cfg: DictDefault, model: PreTrainedModel, safe_serialization: bool
 ):
@@ -282,12 +260,49 @@ def save_trained_model(
             else:
                 state_dict_type = cfg.fsdp_config.state_dict_type
             trainer.accelerator.state.fsdp_plugin.set_state_dict_type(state_dict_type)
-        trainer.save_model(cfg.output_dir)
+        trainer.save_model(cfg.output_dir)  # only handles FULL_STATE_DICT
         if state_dict_type == "SHARDED_STATE_DICT":
             LOG.info(
                 "The final model was saved with a sharded state dict. Please ensure you merge "
                 "the sharded weights with `merge-sharded-fsdp-weights`."
             )
+            checkpoint_dir = determine_last_checkpoint(cfg, update=False)
+            if (
+                not (Path(cfg.output_dir) / "model.safetensors.index.json").exists()
+                and checkpoint_dir
+            ):
+                # import here to prevent circular import
+                from axolotl.cli.merge_sharded_fsdp_weights import merge_fsdp_weights
+
+                fsdp_dir = Path(checkpoint_dir) / "pytorch_model_fsdp_0"
+                merged_path = str(Path(cfg.output_dir) / "merged")
+                merge_fsdp_weights(
+                    checkpoint_dir=str(fsdp_dir),
+                    output_path=merged_path,
+                    safe_serialization=True,
+                )
+                trainer.accelerator.wait_for_everyone()
+                if trainer.accelerator.is_main_process:
+                    # move all files in merged_path to cfg.output_dir
+                    for merged_file in Path(merged_path).iterdir():
+                        shutil.move(str(merged_file), cfg.output_dir)
+                    shutil.rmtree(merged_path)  # remove what should be an empty dir
+        # TODO(wing):see https://github.com/huggingface/transformers/pull/40207
+        # cleanup the FSDP prefix in the model config.json
+        if trainer.accelerator.is_main_process:
+            with open(
+                Path(cfg.output_dir) / "config.json", "r", encoding="utf-8"
+            ) as config_file_io:
+                # read the model config as an OrderedDict
+                config = json.load(config_file_io, object_pairs_hook=OrderedDict)
+                config["architectures"] = [
+                    name.lstrip("FSDP") for name in config["architectures"]
+                ]
+            # write the updated model config back
+            with open(
+                os.path.join(cfg.output_dir, "config.json"), "w", encoding="utf-8"
+            ) as config_file_io:
+                json.dump(config, config_file_io, indent=2)
     elif cfg.deepspeed and is_deepspeed_zero3_enabled():
         # Copied over from: https://github.com/huggingface/accelerate/blob/5ae611118057232f441055f7ef9ba0b0f2b8d533/docs/source/usage_guides/deepspeed.md#saving-and-loading
         trainer.accelerator.wait_for_everyone()
@@ -564,7 +579,7 @@ def train(
     setup_model_card(cfg)
 
     # Execute the training
-    resume_from_checkpoint = determine_resume_checkpoint(cfg)
+    resume_from_checkpoint = determine_last_checkpoint(cfg)
     execute_training(cfg, trainer, resume_from_checkpoint)
 
     # clear cache

src/axolotl/utils/collators/mm_chat.py

@@ -5,7 +5,6 @@ Collators for multi-modal chat messages and packing
 from dataclasses import dataclass
 from typing import Any, Optional, Union
 
-import torch
 from torch import Tensor
 from transformers import PreTrainedTokenizerBase
 from transformers.data.data_collator import DataCollatorMixin
@@ -42,62 +41,19 @@ class MultiModalChatDataCollator(DataCollatorMixin):
         examples = self.processing_strategy(examples)
 
         # Initialize batch
-        batch: dict[str, Any] = {}
+        messages = [ex["messages"] for ex in examples]
 
-        # Process each example
-        for example in examples:
-            # Apply chat template to process the example
-            # This method requires transformers>=4.49.0
-            result = self.processing_strategy.processor.apply_chat_template(
-                example["messages"],
-                add_generation_prompt=False,
-                tokenize=True,
-                return_tensors="pt",
-                padding=True,
-                return_dict=True,
-                chat_template=self.processing_strategy.chat_template,
-            )
-
-            # TODO: Check if need handling for len(input_ids) > sequence_len
-
-            # Add the processed tensors to our batch
-            for key in result.keys():
-                if key not in batch:
-                    batch[key] = []
-
-                batch[key].append(result[key].squeeze(0))
-
-        # Pad sequences to the same length
-        input_ids = torch.nn.utils.rnn.pad_sequence(
-            batch["input_ids"],
-            batch_first=True,
-            padding_value=self.tokenizer.pad_token_id,
+        batch = self.processing_strategy.processor.apply_chat_template(
+            messages,
+            add_generation_prompt=False,
+            tokenize=True,
+            return_tensors="pt",
+            padding=True,
+            return_dict=True,
+            chat_template=self.processing_strategy.chat_template,
         )
 
-        attention_mask = torch.nn.utils.rnn.pad_sequence(
-            batch["attention_mask"], batch_first=True, padding_value=0
-        )
-
-        # Create the final batch
-        final_batch = {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-        }
-
-        for key, val in batch.items():
-            if key in ["input_ids", "attention_mask"]:
-                continue
-
-            if key in ["token_type_ids", "cross_attention_mask"]:
-                final_batch[key] = torch.nn.utils.rnn.pad_sequence(
-                    val, batch_first=True, padding_value=0
-                )
-            else:
-                final_batch[key] = torch.stack(val)
-
         # Process the labels
-        final_batch["labels"] = self.processing_strategy.process_labels(
-            final_batch["input_ids"]
-        )
+        batch["labels"] = self.processing_strategy.process_labels(batch["input_ids"])
 
-        return final_batch
+        return batch

src/axolotl/utils/data/sft.py

@@ -28,7 +28,7 @@ from axolotl.utils.data.shared import (
 )
 from axolotl.utils.data.utils import (
     deduplicate_and_log_datasets,
-    drop_long_seq_in_dataset,
+    handle_long_seq_in_dataset,
     retry_on_request_exceptions,
 )
 from axolotl.utils.data.wrappers import get_dataset_wrapper
@@ -339,9 +339,9 @@ def _load_raw_datasets(
 
     if not cfg.skip_prepare_dataset:
         if split == "test" and cfg.eval_sequence_len:
-            dataset = drop_long_seq_in_dataset(dataset, cfg.eval_sequence_len, cfg)
+            dataset = handle_long_seq_in_dataset(dataset, cfg.eval_sequence_len, cfg)
         else:
-            dataset = drop_long_seq_in_dataset(dataset, cfg.sequence_len, cfg)
+            dataset = handle_long_seq_in_dataset(dataset, cfg.sequence_len, cfg)
         if cfg.sample_packing:
             dataset, _ = process_datasets_for_packing(cfg, dataset, None)
 

src/axolotl/utils/data/utils.py

@@ -148,7 +148,36 @@ def deduplicate_and_log_datasets(
     return dataset, other_dataset
 
 
-def drop_long_seq_in_dataset(
+def truncate_long_seq(sample, sequence_len=2048, min_sequence_len=2):
+    """
+    Truncate samples whose sequence length is too long (> sequence_len)
+    or drop those too short (< min_sequence_len).
+    """
+    min_sequence_len = min_sequence_len or 2
+
+    input_ids = sample["input_ids"]
+    results = []
+
+    # Batched (input_ids is a list of lists)
+    for i, seq in enumerate(input_ids):
+        length = len(seq)
+        if length < min_sequence_len:
+            results.append(False)
+        elif length > sequence_len:
+            sample["input_ids"][i] = seq[:sequence_len]
+            if "attention_mask" in sample:
+                sample["attention_mask"][i] = sample["attention_mask"][i][:sequence_len]
+            if "labels" in sample:
+                sample["labels"][i] = sample["labels"][i][:sequence_len]
+            if "position_ids" in sample:
+                sample["position_ids"][i] = sample["position_ids"][i][:sequence_len]
+            results.append(True)
+        else:
+            results.append(True)
+    return results
+
+
+def handle_long_seq_in_dataset(
     dataset: Dataset, sequence_len: int, cfg: DictDefault
 ) -> Dataset:
     """Remove sequences longer than configured maximum from dataset.
@@ -192,8 +221,21 @@ def drop_long_seq_in_dataset(
     if filter_map_kwargs:
         drop_long_kwargs["desc"] = f"Dropping Long Sequences (>{sequence_len})"
 
+    excess_length_strategy = (cfg.excess_length_strategy or "drop").lower()
+    if excess_length_strategy == "truncate":
+        process_fn = functools.partial(
+            truncate_long_seq,
+            sequence_len=sequence_len,
+            min_sequence_len=cfg.min_sample_len,
+        )
+        drop_long_kwargs["desc"] = (
+            f"Truncating/Filtering Sequences (target_len={sequence_len})"
+        )
+    else:
+        process_fn = drop_long
+
     dataset = dataset.filter(
-        drop_long,
+        process_fn,
         batched=True,
         **filter_map_kwargs,
         **drop_long_kwargs,
@@ -201,6 +243,11 @@ def drop_long_seq_in_dataset(
     if prior_len:
         dropped = prior_len - len(dataset)
         if dropped:
-            LOG.warning(f"Dropped {dropped} long samples from dataset")
+            action = (
+                "truncated/filtered"
+                if excess_length_strategy == "truncate"
+                else "dropped"
+            )
+            LOG.warning(f"{action.title()} {dropped} samples from dataset")
 
     return dataset

src/axolotl/utils/schemas/config.py

@@ -109,6 +109,12 @@ class AxolotlInputConfig(
             "description": "Don't upcast the embeddings to float32 when using PEFT. Useful for low-VRAM GPUs"
         },
     )
+    reinit_weights: bool | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "Reinitialize model weights randomly instead of loading pretrained weights"
+        },
+    )
 
     trainer_cls: str | None = Field(
         default=None,
@@ -414,6 +420,12 @@ class AxolotlInputConfig(
             "description": "The maximum length of an input to train with, this should typically be less than 2048 as most models have a token/context limit of 2048"
         },
     )
+    excess_length_strategy: Literal["drop", "truncate"] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": "What to do when a tokenized row exceeds sequence_len. 'drop' removes the row; 'truncate' slices tensors to sequence_len. Defaults to 'drop' for backward compatibility."
+        },
+    )
     eval_sequence_len: int | None = Field(
         default=None,
         json_schema_extra={

src/axolotl/utils/schemas/validation.py

@@ -3,6 +3,7 @@
 # pylint: disable=too-many-boolean-expressions
 
 import json
+import sys
 import tempfile
 from pathlib import Path
 
@@ -369,10 +370,10 @@ class TrainingValidationMixin:
                 "see speed improvements. Please consider setting `torch_compile: "
                 "true` in your config."
             )
+        fsdp_config = data.get("fsdp_config") or {}
         if data.get("fp8") and (
-            data.get("fsdp_config", {}).get("activation_checkpointing", False) is True
-            or data.get("fsdp_config", {}).get("fsdp_activation_checkpointing", False)
-            is True
+            fsdp_config.get("activation_checkpointing", False) is True
+            or fsdp_config.get("fsdp_activation_checkpointing", False) is True
         ):
             LOG.warning(
                 "FP8 + FSDP2 + activation checkpointing may be slower than BF16 "
@@ -817,13 +818,13 @@ class OptimizationValidationMixin:
     @model_validator(mode="before")
     @classmethod
     def check_fsdp_version_in_fsdp_config(cls, data):
-        if data.get("fsdp_config"):
-            if data.get("fsdp_config", {}).get("fsdp_version"):
-                LOG.warning(
-                    "Configuring `fsdp_version` in `fsdp_config` is deprecated. "
-                    "Please configure `fsdp_version` as a top-level field."
-                )
-                data["fsdp_version"] = data.get("fsdp_config").pop("fsdp_version")
+        fsdp_config = data.get("fsdp_config") or {}
+        if fsdp_config and fsdp_config.get("fsdp_version"):
+            LOG.warning(
+                "Configuring `fsdp_version` in `fsdp_config` is deprecated. "
+                "Please configure `fsdp_version` as a top-level field."
+            )
+            data["fsdp_version"] = fsdp_config.pop("fsdp_version")
         return data
 
     @model_validator(mode="before")
@@ -1151,10 +1152,8 @@ class ModelCompatibilityValidationMixin:
     @classmethod
     def check_gpt_oss_fsdp_loading(cls, data):
         if data.get("model_quantization_config", "") == "Mxfp4Config":
-            if (
-                data.get("fsdp_config", {}).get("cpu_ram_efficient_loading", False)
-                is True
-            ):
+            fsdp_config = data.get("fsdp_config") or {}
+            if fsdp_config.get("cpu_ram_efficient_loading", False) is True:
                 raise ValueError(
                     "FSDP cpu_ram_efficient_loading is not supported for Mxfp4Config model quantization."
                 )
@@ -1251,10 +1250,26 @@ class ComplexValidationMixin:
 
             try:
                 import transformers.modeling_flash_attention_utils
+                from transformers.utils import is_flash_attn_greater_or_equal
 
                 # pylint: disable=protected-access
-                transformers.modeling_flash_attention_utils._flash_supports_window_size = (
-                    transformers.modeling_flash_attention_utils._flash_supports_window
+                transformers.modeling_flash_attention_utils._flash_supports_window = (
+                    True
+                )
+                setattr(
+                    sys.modules["transformers.modeling_flash_attention_utils"],
+                    "_flash_supports_window",
+                    True,
+                )
+                setattr(
+                    sys.modules["transformers.modeling_flash_attention_utils"],
+                    "_flash_supports_window_size",
+                    True,
+                )
+                setattr(
+                    sys.modules["transformers.modeling_flash_attention_utils"],
+                    "is_flash_attn_greater_or_equal",
+                    is_flash_attn_greater_or_equal,
                 )
                 import ring_flash_attn  # noqa: F401 # pylint:disable=unused-import
             except ImportError as exception:

src/axolotl/utils/train.py

@@ -0,0 +1,45 @@
+"""Training utils for checkpoints"""
+
+from pathlib import Path
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
+
+LOG = get_logger(__name__)
+
+
+def determine_last_checkpoint(cfg: DictDefault, update: bool = True) -> str | None:
+    """
+    Determine the checkpoint to resume from based on configuration.
+
+    Args:
+        cfg: Dictionary mapping `axolotl` config keys to values.
+        update: Whether to update the config with the determined checkpoint
+
+    Returns:
+        Path to the checkpoint to resume from, or `None` if not resuming.
+    """
+    last_checkpoint = None
+    checkpoints = sorted(
+        (
+            p
+            for p in Path(cfg.output_dir).glob("checkpoint-*")
+            if p.name.split("-")[-1].isdigit()
+        ),
+        key=lambda p: int(p.name.split("-")[-1]),
+    )
+    if checkpoints:
+        last_checkpoint = str(checkpoints[-1])
+        if not update:
+            return last_checkpoint
+
+    if (
+        cfg.resume_from_checkpoint is None
+        and cfg.auto_resume_from_checkpoints
+        and last_checkpoint is not None
+    ):
+        cfg.resume_from_checkpoint = last_checkpoint
+        LOG.info(
+            f"Using Auto-resume functionality to start with checkpoint at {cfg.resume_from_checkpoint}"
+        )
+    return cfg.resume_from_checkpoint

tests/e2e/patched/test_fsdp2_qlora.py

@@ -1,126 +1,28 @@
-"""Integration tests for FSDP Params4bit patches."""
+"""Integration tests for FSDP2 Params4bit patches."""
 
-from unittest.mock import Mock, patch
-
-import bitsandbytes as bnb
 import pytest
-import torch
 from torch.distributed.fsdp._fully_shard._fsdp_param import FSDPParam
 
-from axolotl.monkeypatch.fsdp2_qlora import (
-    apply_bnb_torch_function_patch,
-    patched_torch_function,
-)
-
-
-@pytest.fixture
-def mock_params4bit():
-    """Create a mock Params4bit instance with test attributes."""
-    mock_instance = Mock()
-    mock_instance.requires_grad = True
-    mock_instance.quant_state = "test_state"
-    mock_instance.blocksize = 128
-    mock_instance.compress_statistics = True
-    mock_instance.quant_type = "fp4"
-    mock_instance.quant_storage = "test_storage"
-    mock_instance.module = "test_module"
-    mock_instance.bnb_quantized = True
-    return mock_instance
-
-
-class TestBnbTorchFunctionPatch:
-    """Test the Params4bit.__torch_function__ patch."""
-
-    def test_apply_patch(self):
-        """Test that the patch can be applied."""
-        with patch("bitsandbytes.nn.modules.Params4bit") as mock_cls:
-            apply_bnb_torch_function_patch()
-            assert hasattr(mock_cls, "__torch_function__")
-            assert isinstance(mock_cls.__torch_function__, classmethod)
-
-    # pylint: disable=redefined-outer-name
-    def test_torch_chunk_preserves_attributes(self, mock_params4bit):
-        """Test that torch.chunk preserves Params4bit attributes."""
-        mock_cls = Mock()
-        chunks = (torch.tensor([1, 2]), torch.tensor([3, 4]))
-
-        with patch("torch.nn.Parameter.__torch_function__", return_value=chunks):
-            result = patched_torch_function(
-                mock_cls,
-                torch.chunk,
-                (type(mock_params4bit),),
-                args=(mock_params4bit, 2),
-            )
-
-            assert isinstance(result, tuple)
-            assert len(result) == 2
-
-            # Check that Params4bit constructor was called with preserved attributes
-            assert mock_cls.call_count == 2
-            for call in mock_cls.call_args_list:
-                kwargs = call[1]
-                assert kwargs["requires_grad"] == mock_params4bit.requires_grad
-                assert kwargs["quant_state"] == mock_params4bit.quant_state
-                assert kwargs["blocksize"] == mock_params4bit.blocksize
-
-    # pylint: disable=redefined-outer-name
-    def test_other_functions_fallback(self, mock_params4bit):
-        """Test that non-chunk/split functions use Parameter fallback."""
-        mock_cls = Mock()
-        fallback_result = torch.tensor([5, 6, 7])
-
-        with patch(
-            "torch.nn.Parameter.__torch_function__", return_value=fallback_result
-        ) as mock_fallback:
-            result = patched_torch_function(
-                mock_cls, torch.add, (type(mock_params4bit),), args=(mock_params4bit, 1)
-            )
-
-            # Should call Parameter.__torch_function__ and return its result
-            mock_fallback.assert_called_once()
-            assert result is fallback_result
-            mock_cls.assert_not_called()
-
 
 class TestFSDPPatchIntegration:
     """Test FSDP patch integration."""
 
     @pytest.mark.integration
-    def test_all_patches_together(self):
+    def test_fsdp2_init_patches(self):
         """Test that all patches can be applied together."""
         from axolotl.monkeypatch.fsdp2_qlora import (
             apply_init_sharded_param_patch,
             apply_init_unsharded_param_patch,
         )
 
-        # Store original methods before patching
-        original_torch_function = getattr(
-            bnb.nn.modules.Params4bit, "__torch_function__", None
-        )
-
         # pylint: disable=protected-access
         original_init_sharded = FSDPParam._init_sharded_param
         original_init_unsharded = FSDPParam.init_unsharded_param
 
         # Apply patches
-        apply_bnb_torch_function_patch()
         apply_init_sharded_param_patch()
         apply_init_unsharded_param_patch()
 
-        # Verify patches were applied
-        current_torch_function = getattr(
-            bnb.nn.modules.Params4bit, "__torch_function__", None
-        )
-        if original_torch_function is not None:
-            assert (
-                current_torch_function != original_torch_function
-            ), "Params4bit.__torch_function__ was not patched"
-        else:
-            assert (
-                current_torch_function is not None
-            ), "Params4bit.__torch_function__ was not added"
-
-        # Check that FSDP methods were patched
         assert (
             # pylint: disable=protected-access
             FSDPParam._init_sharded_param

tests/e2e/test_diffusion.py

@@ -0,0 +1,119 @@
+"""E2E smoke test for diffusion training plugin."""
+
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+from tests.e2e.utils import check_model_output_exists
+
+
+class TestDiffusion:
+    """Test case for diffusion training plugin."""
+
+    def test_diffusion_smoke_test(self, temp_dir):
+        """
+        Smoke test for diffusion training to ensure the plugin loads and trains without
+        error.
+        """
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 256,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "bf16": True,
+                "save_safetensors": True,
+                "save_first_step": False,
+                "logging_steps": 1,
+                "eval_steps": 3,
+                # Diffusion-specific config
+                "plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
+                "diffusion_mask_token_id": 16,
+                "diffusion_eps": 1e-3,
+                "diffusion_importance_weighting": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)
+
+    def test_diffusion_sft_labels(self, temp_dir):
+        """Test that diffusion training properly handles SFT data with labels."""
+        cfg = DictDefault(
+            {
+                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "tokenizer_type": "AutoTokenizer",
+                "trust_remote_code": True,
+                "sequence_len": 256,
+                "val_set_size": 0.1,
+                "special_tokens": {
+                    "pad_token": "<|endoftext|>",
+                },
+                "datasets": [
+                    {
+                        "path": "mhenrichsen/alpaca_2k_test",
+                        "type": "alpaca",
+                    },
+                ],
+                "num_epochs": 1,
+                "max_steps": 3,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 1,
+                "output_dir": temp_dir,
+                "learning_rate": 0.0001,
+                "optimizer": "adamw_torch",
+                "lr_scheduler": "cosine",
+                "bf16": True,
+                "save_safetensors": True,
+                "save_first_step": False,
+                "logging_steps": 1,
+                "eval_steps": 2,
+                # Diffusion-specific config
+                "plugins": ["axolotl.integrations.diffusion.DiffusionPlugin"],
+                "diffusion_mask_token_id": 16,
+                "diffusion_eps": 1e-3,
+                "diffusion_importance_weighting": True,
+                # Ensure we have proper SFT labels
+                "train_on_inputs": False,
+            }
+        )
+
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        dataset_meta = load_datasets(cfg=cfg)
+
+        # Verify that the dataset has labels
+        sample = dataset_meta.train_dataset[0]
+        assert "labels" in sample, "SFT dataset should have labels"
+
+        # Check that some labels are -100 (prompt tokens)
+        labels = sample["labels"]
+        if hasattr(labels, "tolist"):
+            labels = labels.tolist()
+        assert -100 in labels, "SFT dataset should have -100 labels for prompt tokens"
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        check_model_output_exists(temp_dir, cfg)

tests/e2e/utils.py

@@ -147,7 +147,11 @@ def require_hopper(test_case):
 
 
 def check_tensorboard(
-    temp_run_dir: str, tag: str, lt_val: float, assertion_err: str
+    temp_run_dir: str,
+    tag: str,
+    lt_val: float,
+    assertion_err: str,
+    rtol: float = 0.02,
 ) -> None:
     """
     helper function to parse and check tensorboard logs
@@ -157,6 +161,7 @@ def check_tensorboard(
     reader = SummaryReader(event_file)
     df = reader.scalars  # pylint: disable=invalid-name
     df = df[(df.tag == tag)]  # pylint: disable=invalid-name
+    lt_val = (1 + rtol) * lt_val
     if "%s" in assertion_err:
         assert df.value.values[-1] < lt_val, assertion_err % df.value.values[-1]
     else:

tests/integrations/test_diffusion.py

@@ -0,0 +1,290 @@
+"""Tests for diffusion model integration."""
+
+# pylint: disable=redefined-outer-name,protected-access
+
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from axolotl.integrations.diffusion.configuration import LlamaForDiffusionConfig
+from axolotl.integrations.diffusion.models import LlamaForDiffusionLM
+from axolotl.utils.dict import DictDefault
+
+
+@pytest.fixture
+def mock_tokenizer():
+    """Create a mock tokenizer."""
+    tokenizer = Mock()
+    tokenizer.bos_token_id = 1
+    tokenizer.eos_token_id = 2
+    tokenizer.pad_token_id = 0
+    return tokenizer
+
+
+@pytest.fixture
+def diffusion_config():
+    """Create a diffusion config."""
+    return LlamaForDiffusionConfig(
+        mask_token_id=32000,
+        eps=1e-3,
+        importance_weighting=False,
+        sample_packing=False,
+        # Basic llama config fields - smaller for testing
+        vocab_size=1000,
+        hidden_size=256,
+        intermediate_size=512,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+    )
+
+
+@pytest.fixture
+def diffusion_model_instance(mock_tokenizer, diffusion_config):
+    """Create a diffusion model instance for testing methods directly."""
+    # Create a minimal model instance for testing
+    model = object.__new__(LlamaForDiffusionLM)
+    model.config = diffusion_config
+    model._special_token_ids = {0, 1, 2}  # pad, bos, eos
+    model.training = True
+    
+    # Set tokenizer
+    model.set_tokenizer(mock_tokenizer)
+    
+    return model
+
+
+class TestDiffusionModel:
+    """Test the DiffusionModel class."""
+
+    def test_forward_process_basic(self, diffusion_model_instance):
+        """Test basic forward process without labels."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+
+        noisy_batch, masked_indices, p_mask = (
+            diffusion_model_instance._forward_process(input_ids, eps=0.1)
+        )
+
+        # Check shapes
+        assert noisy_batch.shape == input_ids.shape
+        assert masked_indices.shape == input_ids.shape
+        assert p_mask.shape == input_ids.shape
+
+        # Check that special tokens are not masked
+        special_token_positions = (input_ids == 1) | (input_ids == 2) | (input_ids == 0)
+        assert not masked_indices[special_token_positions].any()
+
+        # Check that mask token is applied
+        mask_token_id = diffusion_model_instance.config.mask_token_id
+        masked_positions = masked_indices
+        if masked_positions.any():
+            assert (noisy_batch[masked_positions] == mask_token_id).all()
+
+    def test_forward_process_with_labels(self, diffusion_model_instance):
+        """Test forward process with SFT labels."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        labels = torch.tensor([[-100, -100, 20, 30, 2]], dtype=torch.long)
+
+        noisy_batch, masked_indices, p_mask = (
+            diffusion_model_instance._forward_process(
+                input_ids, labels=labels, eps=0.1
+            )
+        )
+
+        # Check shapes
+        assert noisy_batch.shape == input_ids.shape
+        assert masked_indices.shape == input_ids.shape
+        assert p_mask.shape == input_ids.shape
+
+        # Check that only answer tokens can be masked (where labels != -100)
+        non_answer_mask = labels == -100
+
+        # No masking should occur on non-answer tokens
+        assert not masked_indices[non_answer_mask].any()
+
+        # p_mask should be the same for all positions (sampled timestep),
+        # but masking is only applied to answer tokens
+        assert p_mask.shape == input_ids.shape
+        # Verify that masked_indices respects the answer mask
+        assert not masked_indices[non_answer_mask].any()
+
+    def test_forward_process_with_attention_mask(self, diffusion_model_instance):
+        """Test forward process with attention mask."""
+        input_ids = torch.tensor([[1, 10, 20, 0]], dtype=torch.long)
+        attention_mask = torch.tensor([[1, 1, 1, 0]], dtype=torch.long)
+
+        _, masked_indices, p_mask = diffusion_model_instance._forward_process(
+            input_ids, attention_mask=attention_mask, eps=0.1
+        )
+
+        # Check that padding tokens are not masked
+        padding_positions = attention_mask == 0
+        assert not masked_indices[padding_positions].any()
+        assert (p_mask[padding_positions] == 0).all()
+
+    def test_bidirectional_attention_mask_no_packing(self, diffusion_model_instance):
+        """Test bidirectional attention mask without sample packing."""
+        input_ids = torch.tensor([[1, 10, 20, 2]], dtype=torch.long)
+
+        mask = diffusion_model_instance._create_bidirectional_attention_mask(
+            input_ids
+        )
+
+        # Should be all-to-all attention
+        expected_shape = (1, 1, 4, 4)
+        assert mask.shape == expected_shape
+        assert mask.all()
+
+    def test_bidirectional_attention_mask_with_packing(
+        self, diffusion_model_instance
+    ):
+        """Test bidirectional attention mask with sample packing."""
+        diffusion_model_instance.config.sample_packing = True
+        input_ids = torch.tensor([[1, 10, 20, 30, 40, 2]], dtype=torch.long)
+        # Sample IDs: first sample (1), second sample (2)
+        attention_mask = torch.tensor([[1, 1, 1, 2, 2, 2]], dtype=torch.long)
+
+        mask = diffusion_model_instance._create_bidirectional_attention_mask(
+            input_ids, attention_mask
+        )
+
+        # Check that tokens within same sample can attend to each other
+        # but not across samples
+        assert mask[0, 0, 0, 1].item()  # First sample tokens can attend to each other
+        assert mask[0, 0, 1, 2].item()
+        assert not mask[0, 0, 0, 3].item()  # Can't attend across samples
+        assert not mask[0, 0, 2, 4].item()
+        assert mask[0, 0, 3, 4].item()  # Second sample tokens can attend to each other
+
+    def test_compute_loss_basic(self, diffusion_model_instance):
+        """Test basic loss computation."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        
+        # Create mock data for loss computation
+        vocab_size = 1000
+        seq_len = 5
+        logits = torch.randn(1, seq_len, vocab_size, requires_grad=True)
+        
+        # Create a simple masked indices tensor (mask middle tokens)
+        masked_indices = torch.tensor([[False, True, True, False, False]], dtype=torch.bool)
+        p_mask = torch.tensor([[0.1, 0.5, 0.5, 0.1, 0.1]], dtype=torch.float)
+
+        loss = diffusion_model_instance._compute_diffusion_loss(
+            input_ids=input_ids,
+            logits=logits,
+            masked_indices=masked_indices,
+            p_mask=p_mask,
+        )
+
+        # Check that loss is computed
+        assert isinstance(loss, torch.Tensor)
+        assert loss.requires_grad
+
+    def test_compute_loss_with_labels(self, diffusion_model_instance):
+        """Test loss computation with SFT labels."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        labels = torch.tensor([[-100, -100, 20, 30, 2]], dtype=torch.long)
+        
+        # Create mock data for loss computation
+        vocab_size = 1000
+        seq_len = 5
+        logits = torch.randn(1, seq_len, vocab_size, requires_grad=True)
+        
+        # Create masked indices that only covers answer tokens
+        masked_indices = torch.tensor([[False, False, True, True, False]], dtype=torch.bool)
+        p_mask = torch.tensor([[0.1, 0.1, 0.5, 0.5, 0.1]], dtype=torch.float)
+
+        loss = diffusion_model_instance._compute_diffusion_loss(
+            input_ids=input_ids,
+            labels=labels,
+            logits=logits,
+            masked_indices=masked_indices,
+            p_mask=p_mask,
+        )
+
+        # Check that loss is computed
+        assert isinstance(loss, torch.Tensor)
+        assert loss.requires_grad
+
+    def test_compute_loss_no_masked_tokens(self, diffusion_model_instance):
+        """Test loss computation when no tokens are masked."""
+        input_ids = torch.tensor([[1, 0, 2]], dtype=torch.long)
+        
+        # Create mock data for loss computation
+        vocab_size = 1000
+        seq_len = 3
+        logits = torch.randn(1, seq_len, vocab_size)
+        
+        # No tokens masked
+        masked_indices = torch.tensor([[False, False, False]], dtype=torch.bool)
+        p_mask = torch.tensor([[0.1, 0.1, 0.1]], dtype=torch.float)
+
+        loss = diffusion_model_instance._compute_diffusion_loss(
+            input_ids=input_ids,
+            logits=logits,
+            masked_indices=masked_indices,
+            p_mask=p_mask,
+        )
+
+        # Loss should be zero when no tokens are masked
+        assert loss.item() == 0.0
+        assert loss.requires_grad
+
+    def test_cache_special_token_ids(self, diffusion_model_instance):
+        """Test caching of special token IDs."""
+        # Should cache BOS, EOS, PAD tokens
+        expected_tokens = {0, 1, 2}  # pad, bos, eos
+        assert diffusion_model_instance._special_token_ids == expected_tokens
+
+    def test_cache_special_token_ids_no_tokenizer(self):
+        """Test caching when no tokenizer is available."""
+        # Mock the parent model initialization to avoid loading pretrained weights
+        with patch('transformers.models.llama.modeling_llama.LlamaForCausalLM.__init__'):
+            model = LlamaForDiffusionLM.__new__(LlamaForDiffusionLM)
+            model._cache_special_token_ids(None)
+            assert model._special_token_ids == set()
+
+    def test_forward_training_mode(self, diffusion_model_instance):
+        """Test forward pass in training mode."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        attention_mask = torch.tensor([[1, 1, 1, 1, 1]], dtype=torch.bool)
+        
+        # Mock the parent forward method
+        with patch.object(diffusion_model_instance.__class__.__bases__[1], 'forward') as mock_forward:
+            mock_output = Mock()
+            mock_output.logits = torch.randn(1, 5, 32000)
+            mock_forward.return_value = mock_output
+            
+            # Set training mode
+            diffusion_model_instance.training = True
+            
+            result = diffusion_model_instance.forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                return_dict=True
+            )
+            
+            # Should call parent forward and compute loss
+            assert mock_forward.called
+            assert hasattr(result, 'loss')
+
+    def test_forward_inference_mode(self, diffusion_model_instance):
+        """Test forward pass in inference mode."""
+        input_ids = torch.tensor([[1, 10, 20, 30, 2]], dtype=torch.long)
+        
+        # Mock the parent forward method
+        with patch.object(diffusion_model_instance.__class__.__bases__[1], 'forward') as mock_forward:
+            mock_output = Mock()
+            mock_forward.return_value = mock_output
+            
+            # Set inference mode
+            diffusion_model_instance.training = False
+            
+            result = diffusion_model_instance.forward(
+                input_ids=input_ids,
+                return_dict=True
+            )
+            
+            # Should just call parent forward without diffusion processing
+            assert mock_forward.called
+            assert result == mock_output

tests/test_packed_batch_sampler.py

@@ -8,7 +8,7 @@ from transformers import AutoTokenizer
 from axolotl.datasets import TokenizedPromptDataset
 from axolotl.prompt_strategies.completion import load
 from axolotl.utils.collators import V2BatchSamplerDataCollatorForSeq2Seq
-from axolotl.utils.data.utils import drop_long_seq_in_dataset
+from axolotl.utils.data.utils import handle_long_seq_in_dataset
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 
@@ -70,7 +70,7 @@ class TestBatchedSamplerPacking:
         )
         train_dataset = concatenate_datasets([dataset_wrapper])
 
-        train_dataset = drop_long_seq_in_dataset(train_dataset, cfg.sequence_len, cfg)
+        train_dataset = handle_long_seq_in_dataset(train_dataset, cfg.sequence_len, cfg)
 
         lengths = get_dataset_lengths(train_dataset)
         batch_sampler = MultipackBatchSampler(

tests/utils/test_train.py

@@ -0,0 +1,24 @@
+"""test for train checkpoint utils"""
+
+import os
+
+from axolotl.utils.dict import DictDefault
+from axolotl.utils.train import determine_last_checkpoint
+
+
+def test_determine_last_checkpoint(temp_dir):
+    cfg = DictDefault(
+        output_dir=temp_dir,
+    )
+    for cpt_idx in [1, 9, 10, 20]:
+        os.makedirs(
+            os.path.join(cfg.output_dir, f"checkpoint-{cpt_idx}"), exist_ok=True
+        )
+
+    last_checkpoint = determine_last_checkpoint(cfg, update=False)
+    assert last_checkpoint == os.path.join(cfg.output_dir, "checkpoint-20")
+
+    cfg.resume_from_checkpoint = None
+    cfg.auto_resume_from_checkpoints = True
+    determine_last_checkpoint(cfg, update=True)
+    assert cfg.resume_from_checkpoint == os.path.join(cfg.output_dir, "checkpoint-20")