Remove deprecated configuration files: deleted config.qmd and finetune copy.yml to streamline project structure and eliminate unused resources.

- Introduced DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING constant in utils.py.
- Updated encode_packed_pretraining function to use this constant instead of a hardcoded value.

.github/CONTRIBUTING.md

@@ -57,13 +57,6 @@ We welcome ideas for improvements and new features. To suggest an enhancement, o
 5. Push your branch to your fork on GitHub.
 6. Open a new pull request against the `main` branch of the axolotl repository. Include a clear and concise description of your changes, referencing any related issues.
 
-#### Skipping CI Checks
-
-You can skip certain CI checks by including specific keywords in your commit messages:
-
-- `[skip ci]` or `skip ci` - Skips all CI checks for that commit
-- `[skip-e2e]` or `skip-e2e` - Skips only end-to-end tests while running other CI checks. You may also include this in the title of your PR to disable end-to-end tests for the entire PR.
-
 ## Style Guidelines
 
 ### Code Style

.github/workflows/tests.yml

@@ -188,44 +188,13 @@ 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: >
+    if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' && !github.event.pull_request.draft }}
-      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, gate-skip-e2e]
+    needs: [pre-commit, pytest, pytest-sdist]
 
     strategy:
       fail-fast: false
@@ -271,16 +240,13 @@ jobs:
           modal run cicd.e2e_tests
 
   docker-e2e-tests:
-    if: >
+    if: ${{ github.repository_owner == 'axolotl-ai-cloud' && !github.event.pull_request.draft }}
-      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, gate-skip-e2e, docker-e2e-tests-1st]
+    needs: [pre-commit, pytest, docker-e2e-tests-1st]
 
     strategy:
       fail-fast: false

TODO.md

@@ -0,0 +1,10 @@
+# todo list
+
+- [] Validation of parameters for combinations that won't work
+
+
+
+## things that are known not to work
+
+- FSDP offload and gradient_checkpointing - https://github.com/pytorch/pytorch/issues/82203
+- adamw_bnb_8bit doesn't play well with FSDP offload

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 && \
-    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 "causal_conv1d @ git+https://github.com/Dao-AILab/causal-conv1d.git@main" && \
     python3 -m pip install --no-cache-dir "mamba_ssm @ git+https://github.com/state-spaces/mamba.git@main" && \
     python3 -m pip cache purge
 

docs/multimodal.qmd

@@ -13,13 +13,10 @@ 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
 
@@ -34,7 +31,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 if specified
+chat_template:  # see in next section
 
 # example dataset
 datasets:
@@ -100,16 +97,6 @@ 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}
@@ -156,26 +143,6 @@ 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.
@@ -214,20 +181,6 @@ 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

@@ -1,58 +0,0 @@
-# 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-vl-lora.yaml

@@ -1,58 +0,0 @@
-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,44 +33,13 @@ Note: Memory usage taken from `device_mem_reserved(gib)` from logs.
 
 ### Training 120B
 
-On 8xH100s, make sure you have ~3TB of free disk space. With each checkpoint clocking in at ~720GB, along with the base
+On 8xH100s
-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,7 +20,6 @@ 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

@@ -0,0 +1,7 @@
+# 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/lfm2/lfm2-350m-fft.yaml

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

examples/smolvlm2/README.md

@@ -1,49 +0,0 @@
-# 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

@@ -1,56 +0,0 @@
-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.47.0
+bitsandbytes==0.46.1
 # 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.2
+transformers==4.55.0
 tokenizers>=0.21.1
 accelerate==1.10.0
 datasets==4.0.0

src/axolotl/cli/merge_sharded_fsdp_weights.py

@@ -10,7 +10,6 @@ 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,
@@ -24,7 +23,6 @@ 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__)
 
@@ -145,6 +143,7 @@ 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`")
@@ -181,6 +180,7 @@ 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,32 +195,11 @@ 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=output_path,
+        output_path=str(Path(parsed_cfg.output_dir) / "merged"),
         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/cli/utils/train.py

@@ -67,12 +67,14 @@ def build_command(base_cmd: list[str], options: dict[str, Any]) -> list[str]:
 
 def generate_config_files(config: str, sweep: str | None) -> Iterator[tuple[str, bool]]:
     """
-    Generate list of configuration files to process. Yields a tuple of the configuration file name and a boolean indicating
+    Generate list of configuration files to process.
-    whether this is a group of configurations (i.e., a sweep).
 
     Args:
         config: Base configuration file
         sweep: Sweep configuration file
+
+    Yields:
+        Tuple of configuration file name and whether this is a group of configurations
     """
 
     if not sweep:

src/axolotl/core/builders/causal.py

@@ -424,7 +424,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
     ):
         if training_args.pretraining:
             if (
-                not self.cfg.pretraining_sample_concatenation
+                self.cfg.pretraining_sample_concatenation is False
                 or self.cfg.micro_batch_size > 1
             ):
                 return DataCollatorForSeq2Seq(self.tokenizer, **kwargs)

src/axolotl/core/trainers/__init__.py

@@ -5,6 +5,7 @@
 
 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

@@ -272,20 +272,6 @@ class AxolotlTrainer(
                     num_workers=self.args.dataloader_num_workers,
                     rank=self.args.process_index,
                 )
-
-        if (
-            self.args.accelerator_config is not None
-            and self.args.accelerator_config.split_batches
-            and self.args.accelerator_config.dispatch_batches
-        ):
-            if self.args.sample_packing and self.args.pretraining:
-                if not self.args.eval_sample_packing and not is_training:
-                    dataloader_params["batch_size"] *= self.accelerator.num_processes
-                else:
-                    dataloader_params["batch_size"] = self.accelerator.num_processes
-            elif not self.args.sample_packing and self.args.pretraining:
-                dataloader_params["batch_size"] *= self.accelerator.num_processes
-
         if self.args.sample_packing and (
             (is_training and not self.args.pretraining)
             or (not is_training and self.args.eval_sample_packing is not False)

src/axolotl/core/trainers/mixins/optimizer.py

@@ -185,12 +185,12 @@ class OptimizerMixin(Trainer):
                                 p.data_ptr(): p.numel() for p in module.parameters()
                             }.values()
                         )
-                        LOG.info(f"skipped {module}: {skipped/2**20}M params")
+                        LOG.info(f"skipped {module}: {skipped / 2 ** 20}M params")
                         manager.register_module_override(
                             module, "weight", {"optim_bits": 32}
                         )
                         LOG.debug(f"bitsandbytes: will optimize {module} in fp32")
-                LOG.info(f"skipped: {skipped/2**20}M params")
+                LOG.info(f"skipped: {skipped / 2 ** 20}M params")
 
         if is_sagemaker_mp_enabled():
             self.optimizer = smp.DistributedOptimizer(  # pylint: disable=attribute-defined-outside-init

src/axolotl/loaders/constants.py

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

src/axolotl/loaders/model.py

@@ -25,7 +25,6 @@ from peft import (
 from torch.distributed import DeviceMesh
 from transformers import (
     AutoModelForCausalLM,
-    AutoModelForImageTextToText,
     AutoModelForVision2Seq,
     AwqConfig,
     BitsAndBytesConfig,
@@ -213,7 +212,6 @@ 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."""
@@ -434,8 +432,6 @@ 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"""
@@ -632,16 +628,6 @@ 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:

src/axolotl/loaders/patch_manager.py

@@ -285,10 +285,12 @@ 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,12 +9,73 @@ 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,15 +20,12 @@ 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:  # pylint: disable=duplicate-code
+try:
     from transformers.modeling_flash_attention_utils import _flash_supports_window
 except ImportError:
-    try:
+    from transformers.modeling_flash_attention_utils import (
-        from transformers.modeling_flash_attention_utils import (
+        _flash_supports_window_size as _flash_supports_window,
-            _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,15 +15,12 @@ import torch
 import torch.distributed as dist
 from torch.distributed import DeviceMesh
 
-try:  # pylint: disable=duplicate-code
+try:
     from transformers.modeling_flash_attention_utils import _flash_supports_window
 except ImportError:
-    try:
+    from transformers.modeling_flash_attention_utils import (
-        from transformers.modeling_flash_attention_utils import (
+        _flash_supports_window_size as _flash_supports_window,
-            _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, SmolVLMProcessor, VoxtralProcessor
+from transformers import ProcessorMixin, 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 user message
+            # if the image key exists, add the image to the first 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,34 +179,26 @@ 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 msg_idx, msg_content in enumerate(processed_example["messages"]):
+                for i, content in enumerate(
-                    if first_user_idx is None and msg_content["role"] == "user":
+                    processed_example["messages"][0]["content"]
-                        first_user_idx = msg_idx
+                ):
-                    for i, content in enumerate(
+                    # Usually datasets created with image columns, don't have it in the messages itself
-                        processed_example["messages"][msg_idx]["content"]
+                    if content["type"] == "image" and all(
+                        k not in content for k in ["image", "url", "path", "base64"]
                     ):
-                        # Usually datasets created with image columns, don't have it in the messages itself
+                        ind_to_add = i
-                        if content["type"] == "image" and all(
+                        break
-                            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 and msg_ind_to_add is not None:
+                if ind_to_add is not None:
-                    processed_example["messages"][msg_ind_to_add]["content"][
+                    processed_example["messages"][0]["content"][ind_to_add][
-                        ind_to_add
+                        "image"
-                    ]["image"] = image_value
+                    ] = image_value
                 else:
-                    # if no image type is found, add it to end of the first user message
+                    # if no image type is found, add it to end of the first message
-                    if first_user_idx is None:
+                    processed_example["messages"][0]["content"].append(
-                        first_user_idx = 0
-                    processed_example["messages"][first_user_idx]["content"].append(
                         {
                             "type": "image",
                             "image": image_value,
@@ -403,24 +395,6 @@ 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,
@@ -428,43 +402,32 @@ 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(
-            **processing_kwargs,
+            processor, chat_template, image_size, image_resize_algorithm
         )
     if chat_template_type == "gemma3":
         return Gemma3ProcessingStrategy(
-            **processing_kwargs,
+            processor, chat_template, image_size, image_resize_algorithm
         )
     if chat_template_type == "gemma3n":
         return Gemma3nProcessingStrategy(
-            **processing_kwargs,
+            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
         )
 
     if isinstance(processor, VoxtralProcessor):
         return VoxtralProcessingStrategy(
-            **processing_kwargs,
+            processor, chat_template, image_size, image_resize_algorithm
         )
 
-    if isinstance(processor, SmolVLMProcessor):
+    raise ValueError(f"Unsupported chat template type: {chat_template_type}")
-        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,21 +129,13 @@ 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 hasattr(val, "tolist"):
+                if k in ["pixel_values"]:
-                    out[k] = (
+                    batch[k] = val.tolist()
-                        val.tolist() if k == "pixel_values" else val.squeeze(0).tolist()
-                    )
                 else:
-                    out[k] = val
+                    batch[k] = val.squeeze().tolist()
-            return out
+            return batch
 
         return self.tokenizer.apply_chat_template(
             conversation,
@@ -441,13 +433,10 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
                 tokenized_prompt["attention_mask"] = [1] * len(input_ids)
             else:
                 input_ids = tokenized_res["input_ids"]
-                tokenized_prompt = dict(tokenized_res)
+                tokenized_prompt = tokenized_res
 
             if not self.train_on_inputs:
-                if isinstance(prompt_ids, dict):
+                user_prompt_len = len(prompt_ids)
-                    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/train.py

@@ -4,14 +4,11 @@ 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
@@ -41,7 +38,6 @@ 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:
@@ -50,7 +46,7 @@ except ImportError:
     BetterTransformer = None
 
 if typing.TYPE_CHECKING:
-    from axolotl.core.builders import HFCausalTrainerBuilder, HFRLTrainerBuilder
+    from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFRLTrainerBuilder
 
 LOG = get_logger(__name__)
 
@@ -128,6 +124,32 @@ 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
 ):
@@ -260,49 +282,12 @@ 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)  # only handles FULL_STATE_DICT
+        trainer.save_model(cfg.output_dir)
         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()
@@ -579,7 +564,7 @@ def train(
     setup_model_card(cfg)
 
     # Execute the training
-    resume_from_checkpoint = determine_last_checkpoint(cfg)
+    resume_from_checkpoint = determine_resume_checkpoint(cfg)
     execute_training(cfg, trainer, resume_from_checkpoint)
 
     # clear cache

src/axolotl/utils/collators/mm_chat.py

@@ -5,6 +5,7 @@ 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
@@ -41,19 +42,62 @@ class MultiModalChatDataCollator(DataCollatorMixin):
         examples = self.processing_strategy(examples)
 
         # Initialize batch
-        messages = [ex["messages"] for ex in examples]
+        batch: dict[str, Any] = {}
 
-        batch = self.processing_strategy.processor.apply_chat_template(
+        # Process each example
-            messages,
+        for example in examples:
-            add_generation_prompt=False,
+            # Apply chat template to process the example
-            tokenize=True,
+            # This method requires transformers>=4.49.0
-            return_tensors="pt",
+            result = self.processing_strategy.processor.apply_chat_template(
-            padding=True,
+                example["messages"],
-            return_dict=True,
+                add_generation_prompt=False,
-            chat_template=self.processing_strategy.chat_template,
+                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,
         )
 
-        # Process the labels
+        attention_mask = torch.nn.utils.rnn.pad_sequence(
-        batch["labels"] = self.processing_strategy.process_labels(batch["input_ids"])
+            batch["attention_mask"], batch_first=True, padding_value=0
+        )
 
-        return batch
+        # 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"]
+        )
+
+        return final_batch

src/axolotl/utils/data/pretraining.py

@@ -10,6 +10,7 @@ from torch.utils.data import RandomSampler
 from transformers import PreTrainedTokenizerBase
 
 from axolotl.utils.collators import PretrainingBatchSamplerDataCollatorForSeq2Seq
+from axolotl.utils.data.utils import DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING
 from axolotl.utils.logging import get_logger
 from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
 from axolotl.utils.trainer import process_pretraining_datasets_for_packing
@@ -259,6 +260,15 @@ def encode_packed_pretraining(
         # FIXME using attention mask unpad/pad with trainer and packed pretraining is broken atm
         # workaround by using the position id logic for now in trainer
         drop_attention_mask=multipack_attn,
+        # pass through handling mode from config via ds_wrapper function
+        handling=(
+            getattr(ds_wrapper, "cfg", {}).get(
+                "sequence_len_overflow_handling",
+                getattr(ds_wrapper, "cfg", {}).get(
+                    "excess_token_handling", DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING
+                ),
+            )
+        ),
     )
 
     sampler = MultipackBatchSampler(

src/axolotl/utils/data/rl.py

@@ -122,6 +122,14 @@ def _map_dataset(
     return dataset
 
 
+def drop_long_rl_seq(sample, rl, tokenizer, sequence_len, handling="drop"):
+    """
+    Backward-compatibility wrapper for legacy imports in tests.
+    Delegates to the new predicate.
+    """
+    return _drop_long_sequences(sample, rl, tokenizer, sequence_len)
+
+
 def _drop_long_sequences(
     sample: dict[str, Any], rl: RLType, tokenizer: Any, sequence_len: int
 ) -> bool:
@@ -155,11 +163,51 @@ def _drop_long_sequences(
         len_chosen = len(tokenizer(chosen, add_special_tokens=False)["input_ids"])
         len_rejected = len(tokenizer(rejected, add_special_tokens=False)["input_ids"])
 
-        return (len_prompt + len_chosen) <= sequence_len and (
+        # Truncate first, then drop if still invalid (although truncate should handle it)
-            len_prompt + len_rejected
+        handling_mode = sample.get("sequence_len_overflow_handling", "drop")
-        ) <= sequence_len
+        if handling_mode == "truncate":
+            # If both sequences fit, return sample unchanged
+            if (len_prompt + len_chosen) <= sequence_len and (
+                len_prompt + len_rejected
+            ) <= sequence_len:
+                result = sample
+            else:
+                # Calculate maximum response length that can fit with the prompt
+                max_response_len = sequence_len - len_prompt
 
-    if rl is RLType.KTO:
+                if max_response_len <= 0:
+                    # Prompt itself exceeds sequence length. Cannot truncate responses to fix it.
+                    LOG.warning(
+                        "Prompt length (%s) exceeds sequence length (%s) for DPO-like sample; dropping",
+                        len_prompt,
+                        sequence_len,
+                    )
+                    result = False
+
+                else:
+                    # Truncate the chosen and rejected responses if needed
+                    if len_chosen > max_response_len:
+                        chosen_tokens = tokenizer(chosen, add_special_tokens=False)[
+                            "input_ids"
+                        ][:max_response_len]
+                        sample["chosen"] = tokenizer.decode(
+                            chosen_tokens, skip_special_tokens=True
+                        )
+
+                    if len_rejected > max_response_len:
+                        rejected_tokens = tokenizer(rejected, add_special_tokens=False)[
+                            "input_ids"
+                        ][:max_response_len]
+                        sample["rejected"] = tokenizer.decode(
+                            rejected_tokens, skip_special_tokens=True
+                        )
+                    result = sample
+        else:  # handling == "drop"
+            result = (len_prompt + len_chosen) <= sequence_len and (
+                len_prompt + len_rejected
+            ) <= sequence_len
+
+    elif rl == RLType.KTO:
         if not (sample.get("prompt") and sample.get("completion")):
             raise ValueError("Prompt and completion keys are required for KTO datasets")
 
@@ -171,12 +219,86 @@ def _drop_long_sequences(
             tokenizer(completion, add_special_tokens=False)["input_ids"]
         )
 
-        return (len_prompt + len_completion) <= sequence_len
+        # Truncate first
+        handling_mode = sample.get("sequence_len_overflow_handling", "drop")
+        if handling_mode == "truncate":
+            # If sequence fits, return sample unchanged
+            if (len_prompt + len_completion) <= sequence_len:
+                result = sample
+            else:
+                # Calculate maximum completion length
+                max_completion_len = sequence_len - len_prompt
 
-    if rl is RLType.GRPO:
+                if max_completion_len <= 0:
-        return True
+                    # Prompt itself exceeds sequence length. Cannot truncate completion to fix it.
+                    LOG.warning(
+                        "Prompt length (%s) exceeds sequence length (%s) for KTO sample; dropping",
+                        len_prompt,
+                        sequence_len,
+                    )
+                    result = False
+                else:
+                    # Truncate the completion if needed
+                    if len_completion > max_completion_len:
+                        completion_tokens = tokenizer(
+                            completion, add_special_tokens=False
+                        )["input_ids"][:max_completion_len]
+                        sample["completion"] = tokenizer.decode(
+                            completion_tokens, skip_special_tokens=True
+                        )
+                    result = sample
+        else:  # handling == "drop"
+            result = (len_prompt + len_completion) <= sequence_len
 
-    raise ValueError("Unknown RL type")
+    elif rl == RLType.GRPO:
+        # For GRPO always keep
+        result = True
+    else:
+        raise ValueError("Unknown RL type")
+
+    return bool(result)
+
+
+def load_prepare_preference_datasets(cfg):
+    def _is_rl_seq_within_sequence_len(sample, rl, tokenizer, sequence_len):
+        """
+        Boolean predicate to check whether a preference-learning sample fits within sequence_len.
+        Used with dataset.filter() after truncation to drop unsalvageable samples.
+        """
+        if rl in (RLType.DPO, RLType.IPO, RLType.ORPO, RLType.SIMPO):
+            if not (
+                sample.get("prompt") and sample.get("chosen") and sample.get("rejected")
+            ):
+                return False
+            prompt = sample["prompt"]
+            chosen = sample["chosen"]
+            rejected = sample["rejected"]
+            len_prompt = len(tokenizer(prompt, add_special_tokens=False)["input_ids"])
+            len_chosen = len(tokenizer(chosen, add_special_tokens=False)["input_ids"])
+            len_rejected = len(
+                tokenizer(rejected, add_special_tokens=False)["input_ids"]
+            )
+            return (len_prompt + len_chosen) <= sequence_len and (
+                len_prompt + len_rejected
+            ) <= sequence_len
+        if rl == RLType.KTO:
+            if not (sample.get("prompt") and sample.get("completion")):
+                return False
+            prompt = sample["prompt"]
+            completion = sample["completion"]
+            len_prompt = len(tokenizer(prompt, add_special_tokens=False)["input_ids"])
+            len_completion = len(
+                tokenizer(completion, add_special_tokens=False)["input_ids"]
+            )
+            return (len_prompt + len_completion) <= sequence_len
+        if rl == RLType.GRPO:
+            # GRPO does not enforce this check here
+            return True
+        return False
+
+    # Legacy shim preserved for backward compatibility; no-op in new flow
+    def load_split(dataset_cfgs, _cfg):  # noqa: F811
+        return None
 
 
 def _load_split(cfg: DictDefault, split: Literal["train", "test"]) -> Dataset:

src/axolotl/utils/data/utils.py

@@ -15,10 +15,12 @@ from datasets import Dataset, IterableDataset
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.logging import get_logger
 from axolotl.utils.samplers.utils import get_dataset_lengths
-from axolotl.utils.trainer import drop_long_seq
+from axolotl.utils.trainer import truncate_or_drop_long_seq
 
 LOG = get_logger(__name__)
 
+DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING = "drop"
+
 
 class RetryStrategy(Enum):
     """Enum for retry strategies."""
@@ -168,10 +170,19 @@ def drop_long_seq_in_dataset(
         )
         return dataset
 
-    drop_long = functools.partial(
+    # Get the handling method from config, default to "drop" for backward compatibility.
-        drop_long_seq,
+    # Support legacy alias "excess_token_handling" as well.
+    handling = cfg.get(
+        "sequence_len_overflow_handling",
+        cfg.get("excess_token_handling", DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING),
+    )
+
+    # Use the function with the specified handling mode
+    seq_handler = functools.partial(
+        truncate_or_drop_long_seq,
         sequence_len=sequence_len,
         min_sequence_len=cfg.min_sample_len,
+        handling=handling,
     )
 
     with contextlib.suppress(AttributeError):
@@ -190,17 +201,31 @@ def drop_long_seq_in_dataset(
 
     drop_long_kwargs = {}
     if filter_map_kwargs:
-        drop_long_kwargs["desc"] = f"Dropping Long Sequences (>{sequence_len})"
+        if handling == "truncate":
+            drop_long_kwargs["desc"] = "Truncating Long Sequences"
+        else:  # handling == "drop"
+            drop_long_kwargs["desc"] = f"Dropping Long Sequences (>{sequence_len})"
 
-    dataset = dataset.filter(
+    if handling == "truncate":
-        drop_long,
+        # Use map for truncate mode
-        batched=True,
+        dataset = dataset.map(
-        **filter_map_kwargs,
+            seq_handler,
-        **drop_long_kwargs,
+            batched=True,
-    )
+            **filter_map_kwargs,
-    if prior_len:
+            **drop_long_kwargs,
-        dropped = prior_len - len(dataset)
+        )
-        if dropped:
+        LOG.info(f"Truncated long samples in dataset to {sequence_len} tokens")
-            LOG.warning(f"Dropped {dropped} long samples from dataset")
+    else:  # handling == "drop"
+        # Use filter for drop mode
+        dataset = dataset.filter(
+            seq_handler,
+            batched=True,
+            **filter_map_kwargs,
+            **drop_long_kwargs,
+        )
+        if prior_len:
+            dropped = prior_len - len(dataset)
+            if dropped:
+                LOG.warning(f"Dropped {dropped} long samples from dataset")
 
     return dataset

src/axolotl/utils/schemas/config.py

@@ -414,6 +414,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"
         },
     )
+    sequence_len_overflow_handling: Literal["drop", "truncate"] = Field(
+        default="drop",
+        json_schema_extra={
+            "description": "How to handle sequences that overflow the sequence_len: 'drop' (remove the sample) or 'truncate' (cut off excess tokens)."
+        },
+    )
     eval_sequence_len: int | None = Field(
         default=None,
         json_schema_extra={

src/axolotl/utils/schemas/validation.py

@@ -3,7 +3,6 @@
 # pylint: disable=too-many-boolean-expressions
 
 import json
-import sys
 import tempfile
 from pathlib import Path
 
@@ -370,10 +369,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 (
-            fsdp_config.get("activation_checkpointing", False) is True
+            data.get("fsdp_config", {}).get("activation_checkpointing", False) is True
-            or fsdp_config.get("fsdp_activation_checkpointing", False) is True
+            or data.get("fsdp_config", {}).get("fsdp_activation_checkpointing", False)
+            is True
         ):
             LOG.warning(
                 "FP8 + FSDP2 + activation checkpointing may be slower than BF16 "
@@ -818,13 +817,13 @@ class OptimizationValidationMixin:
     @model_validator(mode="before")
     @classmethod
     def check_fsdp_version_in_fsdp_config(cls, data):
-        fsdp_config = data.get("fsdp_config") or {}
+        if data.get("fsdp_config"):
-        if fsdp_config and fsdp_config.get("fsdp_version"):
+            if data.get("fsdp_config", {}).get("fsdp_version"):
-            LOG.warning(
+                LOG.warning(
-                "Configuring `fsdp_version` in `fsdp_config` is deprecated. "
+                    "Configuring `fsdp_version` in `fsdp_config` is deprecated. "
-                "Please configure `fsdp_version` as a top-level field."
+                    "Please configure `fsdp_version` as a top-level field."
-            )
+                )
-            data["fsdp_version"] = fsdp_config.pop("fsdp_version")
+                data["fsdp_version"] = data.get("fsdp_config").pop("fsdp_version")
         return data
 
     @model_validator(mode="before")
@@ -1152,8 +1151,10 @@ class ModelCompatibilityValidationMixin:
     @classmethod
     def check_gpt_oss_fsdp_loading(cls, data):
         if data.get("model_quantization_config", "") == "Mxfp4Config":
-            fsdp_config = data.get("fsdp_config") or {}
+            if (
-            if fsdp_config.get("cpu_ram_efficient_loading", False) is True:
+                data.get("fsdp_config", {}).get("cpu_ram_efficient_loading", False)
+                is True
+            ):
                 raise ValueError(
                     "FSDP cpu_ram_efficient_loading is not supported for Mxfp4Config model quantization."
                 )
@@ -1250,26 +1251,10 @@ 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 = (
+                transformers.modeling_flash_attention_utils._flash_supports_window_size = (
-                    True
+                    transformers.modeling_flash_attention_utils._flash_supports_window
-                )
-                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

@@ -1,45 +0,0 @@
-"""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

src/axolotl/utils/trainer.py

@@ -233,6 +233,114 @@ def drop_long_seq(sample, sequence_len=2048, min_sequence_len=2):
     return results
 
 
+def truncate_or_drop_long_seq(
+    sample, sequence_len=2048, min_sequence_len=2, handling="drop"
+):
+    """
+    Either drop or truncate samples whose sequence length is either too long (> sequence_len)
+    or too short (< min_sequence_len).
+
+    If handling is "drop":
+        - Samples that are too short or too long will be dropped
+    If handling is "truncate":
+        - Samples that are too short will still be dropped
+        - Samples that are too long will be truncated to sequence_len
+
+    Works for both single-example (list[int]) or batched (list[list[int]]).
+    Returns either a boolean/list of booleans (for drop mode) or the modified sample (for truncate mode).
+    """
+    min_sequence_len = min_sequence_len or 2
+    result = None
+
+    if handling == "drop":
+        return drop_long_seq(sample, sequence_len, min_sequence_len)
+
+    input_ids = sample["input_ids"]
+
+    # Edge case: if input_ids is empty
+    if not input_ids:
+        result = False if handling == "drop" else sample
+    # Single example (input_ids is a list of int)
+    elif isinstance(input_ids[0], int):
+        length = len(input_ids)
+
+        # Handle samples that are too short - always drop them
+        if length < min_sequence_len:
+            result = False if handling == "drop" else sample
+        # If truncation is enabled and the sample is too long, truncate it
+        elif length > sequence_len and handling == "truncate":
+            sample["input_ids"] = input_ids[:sequence_len]
+
+            # Also truncate attention_mask if present
+            if "attention_mask" in sample:
+                sample["attention_mask"] = sample["attention_mask"][:sequence_len]
+
+            # Also truncate labels if present
+            if "labels" in sample:
+                sample["labels"] = sample["labels"][:sequence_len]
+
+            # Also truncate position_ids if present
+            if "position_ids" in sample:
+                sample["position_ids"] = sample["position_ids"][:sequence_len]
+
+            # Update length if present
+            if "length" in sample:
+                sample["length"] = sequence_len
+
+            result = sample
+        # For drop mode or if the sample doesn't exceed max length
+        else:
+            result = (
+                min_sequence_len <= length <= sequence_len
+                if handling == "drop"
+                else sample
+            )
+    # Batched (input_ids is a list of lists)
+    else:
+        if handling == "drop":
+            results = []
+            for seq in input_ids:
+                length = len(seq)
+                results.append(min_sequence_len <= length <= sequence_len)
+            result = results
+        else:  # truncate
+            # Check each sequence in the batch
+            for i, seq in enumerate(input_ids):
+                length = len(seq)
+
+                # Skip sequences that are too short
+                if length < min_sequence_len:
+                    continue
+
+                # Truncate sequences that are too long
+                if length > sequence_len:
+                    input_ids[i] = seq[:sequence_len]
+
+                    # Also truncate attention_mask if present
+                    if "attention_mask" in sample:
+                        sample["attention_mask"][i] = sample["attention_mask"][i][
+                            :sequence_len
+                        ]
+
+                    # Also truncate labels if present
+                    if "labels" in sample:
+                        sample["labels"][i] = sample["labels"][i][:sequence_len]
+
+                    # Also truncate position_ids if present
+                    if "position_ids" in sample:
+                        sample["position_ids"][i] = sample["position_ids"][i][
+                            :sequence_len
+                        ]
+
+                    # Update length if present
+                    if "length" in sample:
+                        sample["length"][i] = sequence_len
+
+            result = sample
+
+    return result
+
+
 def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
     drop_attn_mask = cfg.model_config_type in ["mamba", "gemma3"]
     if drop_attn_mask:
@@ -368,15 +476,33 @@ def process_datasets_for_packing(cfg, train_dataset, eval_dataset):
 
 
 def process_pretraining_datasets_for_packing(
-    train_dataset, sequence_len, skip_position_ids=True, drop_attention_mask=False
+    train_dataset,
+    sequence_len,
+    skip_position_ids=True,
+    drop_attention_mask=False,
+    handling="drop",
 ):
-    drop_long = partial(drop_long_seq, sequence_len=sequence_len)
+    # Define the function to use for handling sequences based on the mode
-
+    seq_handler_fn = partial(
-    train_dataset = train_dataset.filter(
+        truncate_or_drop_long_seq,
-        drop_long,
+        sequence_len=sequence_len,
-        desc="Dropping Long Sequences",
+        handling=handling,  # Pass handling mode
-        load_from_cache_file=False,
     )
+
+    # Use map for truncate mode and filter for drop mode
+    if handling == "truncate":
+        train_dataset = train_dataset.map(
+            seq_handler_fn,
+            desc="Truncating Long Sequences",
+            load_from_cache_file=False,
+        )
+    else:  # handling == "drop"
+        train_dataset = train_dataset.filter(
+            seq_handler_fn,  # Use the same function, it returns boolean for drop mode
+            desc="Dropping Long Sequences",
+            load_from_cache_file=False,
+        )
+
     if not skip_position_ids:
         train_dataset = train_dataset.map(
             add_position_ids,

tests/e2e/patched/test_fsdp2_qlora.py

@@ -1,28 +1,126 @@
-"""Integration tests for FSDP2 Params4bit patches."""
+"""Integration tests for FSDP 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_fsdp2_init_patches(self):
+    def test_all_patches_together(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/utils.py

@@ -147,11 +147,7 @@ def require_hopper(test_case):
 
 
 def check_tensorboard(
-    temp_run_dir: str,
+    temp_run_dir: str, tag: str, lt_val: float, assertion_err: str
-    tag: str,
-    lt_val: float,
-    assertion_err: str,
-    rtol: float = 0.02,
 ) -> None:
     """
     helper function to parse and check tensorboard logs
@@ -161,7 +157,6 @@ 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/test_data.py

@@ -3,10 +3,12 @@ test module for the axolotl.utils.data module
 """
 
 import unittest
+from unittest.mock import MagicMock
 
 from transformers import LlamaTokenizer
 
 from axolotl.utils.data import encode_pretraining, md5
+from axolotl.utils.data.rl import drop_long_rl_seq
 
 from tests.hf_offline_utils import enable_hf_offline
 
@@ -64,5 +66,254 @@ class TestEncodePretraining(unittest.TestCase):
         )
 
 
+class TestDropLongRLSeq(unittest.TestCase):
+    """
+    Tests for the drop_long_rl_seq function.
+    """
+
+    def setUp(self):
+        # Mock tokenizer that returns length based on input string length
+        self.tokenizer = MagicMock()
+
+        def side_effect_func(
+            text, add_special_tokens=False
+        ):  # pylint: disable=unused-argument
+            return {"input_ids": list(range(len(text)))}
+
+        self.tokenizer.side_effect = side_effect_func
+        self.tokenizer.decode = lambda tokens, skip_special_tokens: "".join(
+            ["x"] * len(tokens)
+        )  # pylint: disable=unused-argument
+
+        self.sequence_len = 20
+
+    def test_dpo_drop_mode_valid(self):
+        """Test DPO drop mode with a valid sample."""
+        sample = {
+            "prompt": "p" * 5,
+            "chosen": "c" * 7,
+            "rejected": "r" * 6,
+        }  # 5+7=12 <= 20, 5+6=11 <= 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="drop"
+        )
+        self.assertTrue(result)
+
+    def test_dpo_drop_mode_invalid_chosen(self):
+        """Test DPO drop mode with chosen too long."""
+        sample = {
+            "prompt": "p" * 5,
+            "chosen": "c" * 16,
+            "rejected": "r" * 6,
+        }  # 5+16=21 > 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="drop"
+        )
+        self.assertFalse(result)
+
+    def test_dpo_drop_mode_invalid_rejected(self):
+        """Test DPO drop mode with rejected too long."""
+        sample = {
+            "prompt": "p" * 5,
+            "chosen": "c" * 7,
+            "rejected": "r" * 16,
+        }  # 5+16=21 > 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="drop"
+        )
+        self.assertFalse(result)
+
+    def test_dpo_truncate_mode_no_truncation_needed(self):
+        """Test DPO truncate mode when no truncation is needed."""
+        sample = {
+            "prompt": "p" * 5,
+            "chosen": "c" * 7,
+            "rejected": "r" * 6,
+        }  # 5+7=12 <= 20, 5+6=11 <= 20
+        original_sample = sample.copy()
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(
+            result, original_sample
+        )  # Should return the original sample unchanged
+
+    def test_dpo_truncate_mode_prompt_too_long(self):
+        """Test DPO truncate mode when the prompt itself is too long."""
+        sample = {"prompt": "p" * 25, "chosen": "c" * 7, "rejected": "r" * 6}
+        original_sample = sample.copy()
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        # Even though truncation isn't possible, the function should return the original sample
+        # for the map operation, assuming downstream filtering will catch it.
+        self.assertEqual(result, original_sample)
+
+    def test_dpo_truncate_mode_chosen_truncated(self):
+        """Test DPO truncate mode when only 'chosen' needs truncation."""
+        prompt_len = 5
+        max_resp_len = self.sequence_len - prompt_len  # 20 - 5 = 15
+        sample = {
+            "prompt": "p" * prompt_len,
+            "chosen": "c" * 18,
+            "rejected": "r" * 10,
+        }  # 5+18=23 > 20, 5+10=15 <= 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(len(result["prompt"]), prompt_len)
+        self.assertEqual(len(result["chosen"]), max_resp_len)  # Truncated to 15
+        self.assertEqual(
+            result["chosen"], "x" * max_resp_len
+        )  # Check decoded truncated value
+        self.assertEqual(len(result["rejected"]), 10)  # Unchanged
+
+    def test_dpo_truncate_mode_rejected_truncated(self):
+        """Test DPO truncate mode when only 'rejected' needs truncation."""
+        prompt_len = 5
+        max_resp_len = self.sequence_len - prompt_len  # 15
+        sample = {
+            "prompt": "p" * prompt_len,
+            "chosen": "c" * 10,
+            "rejected": "r" * 18,
+        }  # 5+10=15 <= 20, 5+18=23 > 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(len(result["prompt"]), prompt_len)
+        self.assertEqual(len(result["chosen"]), 10)  # Unchanged
+        self.assertEqual(len(result["rejected"]), max_resp_len)  # Truncated to 15
+        self.assertEqual(
+            result["rejected"], "x" * max_resp_len
+        )  # Check decoded truncated value
+
+    def test_dpo_truncate_mode_both_truncated(self):
+        """Test DPO truncate mode when both 'chosen' and 'rejected' need truncation."""
+        prompt_len = 8
+        max_resp_len = self.sequence_len - prompt_len  # 20 - 8 = 12
+        sample = {
+            "prompt": "p" * prompt_len,
+            "chosen": "c" * 15,
+            "rejected": "r" * 14,
+        }  # 8+15=23 > 20, 8+14=22 > 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(len(result["prompt"]), prompt_len)
+        self.assertEqual(len(result["chosen"]), max_resp_len)  # Truncated to 12
+        self.assertEqual(result["chosen"], "x" * max_resp_len)
+        self.assertEqual(len(result["rejected"]), max_resp_len)  # Truncated to 12
+        self.assertEqual(result["rejected"], "x" * max_resp_len)
+
+    def test_dpo_truncate_mode_no_truncation_needed_but_long(self):
+        """Test DPO truncate mode where individual parts fit but combined don't, but no truncation happens."""
+        # This tests the case where len(chosen) <= max_resp_len and len(rejected) <= max_resp_len
+        # but the initial check failed because e.g. prompt + chosen > sequence_len
+        # The current logic *will* truncate if len(chosen) > max_resp_len.
+        # Let's test a case where one is slightly too long causing the initial fail,
+        # but the other fits *within* the max_response_len, so only one gets truncated.
+        prompt_len = 10
+        max_resp_len = self.sequence_len - prompt_len  # 10
+        sample = {
+            "prompt": "p" * prompt_len,
+            "chosen": "c" * 11,
+            "rejected": "r" * 9,
+        }  # 10+11=21 > 20, 10+9=19 <= 20
+        result = drop_long_rl_seq(
+            sample, "dpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(len(result["prompt"]), prompt_len)
+        self.assertEqual(len(result["chosen"]), max_resp_len)  # Truncated to 10
+        self.assertEqual(result["chosen"], "x" * max_resp_len)
+        self.assertEqual(len(result["rejected"]), 9)  # Unchanged, as 9 <= 10
+
+    # Add similar tests for KTO if needed, checking prompt + completion length
+
+    def test_kto_drop_mode_valid(self):
+        """Test KTO drop mode with a valid sample."""
+        sample = {"prompt": "p" * 5, "completion": "c" * 14}  # 5+14=19 <= 20
+        result = drop_long_rl_seq(
+            sample, "kto", self.tokenizer, self.sequence_len, handling="drop"
+        )
+        self.assertTrue(result)
+
+    def test_kto_drop_mode_invalid(self):
+        """Test KTO drop mode with an invalid sample."""
+        sample = {"prompt": "p" * 5, "completion": "c" * 16}  # 5+16=21 > 20
+        result = drop_long_rl_seq(
+            sample, "kto", self.tokenizer, self.sequence_len, handling="drop"
+        )
+        self.assertFalse(result)
+
+    def test_kto_truncate_mode_no_truncation_needed(self):
+        """Test KTO truncate mode when no truncation is needed."""
+        sample = {"prompt": "p" * 5, "completion": "c" * 14}  # 5+14=19 <= 20
+        original_sample = sample.copy()
+        result = drop_long_rl_seq(
+            sample, "kto", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(result, original_sample)
+
+    def test_kto_truncate_mode_prompt_too_long(self):
+        """Test KTO truncate mode when the prompt itself is too long."""
+        sample = {"prompt": "p" * 25, "completion": "c" * 7}
+        original_sample = sample.copy()
+        result = drop_long_rl_seq(
+            sample, "kto", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(result, original_sample)  # Returns original sample
+
+    def test_kto_truncate_mode_completion_truncated(self):
+        """Test KTO truncate mode when completion needs truncation."""
+        prompt_len = 8
+        max_comp_len = self.sequence_len - prompt_len  # 20 - 8 = 12
+        sample = {"prompt": "p" * prompt_len, "completion": "c" * 15}  # 8+15=23 > 20
+        result = drop_long_rl_seq(
+            sample, "kto", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(len(result["prompt"]), prompt_len)
+        self.assertEqual(len(result["completion"]), max_comp_len)  # Truncated to 12
+        self.assertEqual(result["completion"], "x" * max_comp_len)
+
+    def test_missing_keys_dpo(self):
+        """Test ValueError raised if keys missing for DPO."""
+        sample = {"prompt": "p"}
+        with self.assertRaisesRegex(
+            ValueError, "Prompt, chosen and rejected keys are required"
+        ):
+            drop_long_rl_seq(sample, "dpo", self.tokenizer, self.sequence_len)
+
+    def test_missing_keys_kto(self):
+        """Test ValueError raised if keys missing for KTO."""
+        sample = {"prompt": "p"}
+        with self.assertRaisesRegex(
+            ValueError, "Prompt and completion keys are required"
+        ):
+            drop_long_rl_seq(sample, "kto", self.tokenizer, self.sequence_len)
+
+    def test_unknown_rl_type(self):
+        """Test ValueError raised for unknown RL type."""
+        sample = {}
+        with self.assertRaisesRegex(ValueError, "Unknown RL type"):
+            drop_long_rl_seq(sample, "xyz", self.tokenizer, self.sequence_len)
+
+    # GRPO test - current implementation always passes
+    def test_grpo_drop(self):
+        """Test GRPO drop mode (currently always True)."""
+        sample = {}
+        result = drop_long_rl_seq(
+            sample, "grpo", self.tokenizer, self.sequence_len, handling="drop"
+        )
+        self.assertTrue(result)
+
+    def test_grpo_truncate(self):
+        """Test GRPO truncate mode (currently returns original sample)."""
+        sample = {"a": 1}
+        result = drop_long_rl_seq(
+            sample, "grpo", self.tokenizer, self.sequence_len, handling="truncate"
+        )
+        self.assertEqual(result, sample)
+
+
 if __name__ == "__main__":
     unittest.main()

tests/test_trainer_utils.py

@@ -0,0 +1,153 @@
+"""Module containing tests for trainer utility functions."""
+
+import unittest
+from functools import partial
+
+from axolotl.utils.trainer import truncate_or_drop_long_seq
+
+
+# Test cases for truncate_or_drop_long_seq
+class TestTruncateOrDropLongSeq(unittest.TestCase):
+    """
+    Test suite for truncate_or_drop_long_seq function.
+    """
+
+    def setUp(self):
+        # Example sequence length settings
+        self.sequence_len = 10
+        self.min_sequence_len = 3
+
+    def test_drop_mode_single(self):
+        """Test drop mode with single examples."""
+        handler = partial(
+            truncate_or_drop_long_seq,
+            sequence_len=self.sequence_len,
+            min_sequence_len=self.min_sequence_len,
+            handling="drop",
+        )
+
+        # Too short
+        sample_short = {"input_ids": [1, 2]}
+        self.assertFalse(handler(sample_short))
+
+        # Too long
+        sample_long = {"input_ids": list(range(self.sequence_len + 1))}
+        self.assertFalse(handler(sample_long))
+
+        # Just right
+        sample_ok = {"input_ids": list(range(self.min_sequence_len))}
+        self.assertTrue(handler(sample_ok))
+
+        # Empty
+        sample_empty = {"input_ids": []}
+        self.assertFalse(handler(sample_empty))
+
+    def test_truncate_mode_single(self):
+        """Test truncate mode with single examples."""
+        handler = partial(
+            truncate_or_drop_long_seq,
+            sequence_len=self.sequence_len,
+            min_sequence_len=self.min_sequence_len,
+            handling="truncate",
+        )
+
+        # Too short (should still be dropped implicitly by filter/map logic upstream,
+        # but the function itself might return the sample or False based on impl.)
+        # Current impl returns the original sample for map if too short, assuming upstream filters.
+        # Let's refine this test - the function *itself* returns the sample if too short when truncating.
+        sample_short = {"input_ids": [1, 2], "labels": [1, 2]}
+        result_short = handler(sample_short)
+        self.assertEqual(result_short["input_ids"], [1, 2])  # Unchanged
+
+        # Too long
+        original_long = list(range(self.sequence_len + 5))
+        sample_long = {"input_ids": list(original_long), "labels": list(original_long)}
+        result_long = handler(sample_long)
+        self.assertEqual(len(result_long["input_ids"]), self.sequence_len)
+        self.assertEqual(result_long["input_ids"], list(range(self.sequence_len)))
+        self.assertEqual(len(result_long["labels"]), self.sequence_len)
+        self.assertEqual(result_long["labels"], list(range(self.sequence_len)))
+
+        # Just right
+        sample_ok = {
+            "input_ids": list(range(self.min_sequence_len)),
+            "labels": list(range(self.min_sequence_len)),
+        }
+        result_ok = handler(sample_ok)
+        self.assertEqual(len(result_ok["input_ids"]), self.min_sequence_len)
+        self.assertEqual(result_ok, sample_ok)  # Should be unchanged
+
+        # Empty
+        sample_empty = {"input_ids": [], "labels": []}
+        result_empty = handler(sample_empty)
+        self.assertEqual(result_empty, sample_empty)  # Unchanged
+
+    def test_drop_mode_batched(self):
+        """Test drop mode with batched examples."""
+        handler = partial(
+            truncate_or_drop_long_seq,
+            sequence_len=self.sequence_len,
+            min_sequence_len=self.min_sequence_len,
+            handling="drop",
+        )
+        sample = {
+            "input_ids": [
+                [1, 2],  # Too short
+                list(range(self.sequence_len + 1)),  # Too long
+                list(range(self.sequence_len)),  # OK (len = 10)
+                list(range(self.min_sequence_len)),  # OK (len = 3)
+                [],  # Empty
+            ]
+        }
+        expected = [False, False, True, True, False]
+        self.assertEqual(handler(sample), expected)
+
+    def test_truncate_mode_batched(self):
+        """Test truncate mode with batched examples."""
+        handler = partial(
+            truncate_or_drop_long_seq,
+            sequence_len=self.sequence_len,
+            min_sequence_len=self.min_sequence_len,
+            handling="truncate",
+        )
+        sample = {
+            "input_ids": [
+                [1, 2],  # Too short
+                list(range(self.sequence_len + 5)),  # Too long
+                list(range(self.sequence_len)),  # OK
+                list(range(self.min_sequence_len)),  # OK
+                [],  # Empty
+            ],
+            "labels": [  # Add labels to test truncation
+                [1, 2],
+                list(range(self.sequence_len + 5)),
+                list(range(self.sequence_len)),
+                list(range(self.min_sequence_len)),
+                [],
+            ],
+        }
+
+        result = handler(sample)
+
+        # Expected results after truncation (too short and empty remain unchanged by this function)
+        expected_input_ids = [
+            [1, 2],  # Unchanged (too short)
+            list(range(self.sequence_len)),  # Truncated
+            list(range(self.sequence_len)),  # Unchanged (OK)
+            list(range(self.min_sequence_len)),  # Unchanged (OK)
+            [],  # Unchanged (Empty)
+        ]
+        expected_labels = [
+            [1, 2],  # Unchanged (too short)
+            list(range(self.sequence_len)),  # Truncated
+            list(range(self.sequence_len)),  # Unchanged (OK)
+            list(range(self.min_sequence_len)),  # Unchanged (OK)
+            [],  # Unchanged (Empty)
+        ]
+
+        self.assertEqual(result["input_ids"], expected_input_ids)
+        self.assertEqual(result["labels"], expected_labels)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/utils/test_train.py

@@ -1,24 +0,0 @@
-"""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")