improve handling and error if fa3 requested but not installeD

* don't set peft_config on grpo to prevent double peft wrap

* remove overrides needed to support bug

* fix grpo tests

* require more CPU for multigpu to help with torch compile for vllm

.github/workflows/base.yml

@@ -47,11 +47,18 @@ jobs:
             pytorch: 2.7.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
           - cuda: "128"
-            cuda_version: 12.6.3
+            cuda_version: 12.8.1
             cudnn_version: ""
             python_version: "3.11"
             pytorch: 2.7.0
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
+          - cuda: "126"
+            cuda_version: 12.6.3
+            cudnn_version: ""
+            python_version: "3.11"
+            pytorch: 2.6.0
+            suffix: "-hopper"
+            torch_cuda_arch_list: "9.0+PTX"
           - cuda: "128"
             cuda_version: 12.8.1
             cudnn_version: ""
@@ -87,7 +94,7 @@ jobs:
           context: .
           file: ${{ matrix.pytorch == 'nightly' && './docker/Dockerfile-base-nightly' || matrix.pytorch == 'next' && './docker/Dockerfile-base-next' || './docker/Dockerfile-base' }}
           push: ${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+          tags: ${{ steps.metadata.outputs.tags }}-base-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}${{ matrix.suffix || '' }}
           labels: ${{ steps.metadata.outputs.labels }}
           build-args: |
             CUDA_VERSION=${{ matrix.cuda_version }}

.github/workflows/main.yml

@@ -31,6 +31,11 @@ jobs:
             python_version: "3.11"
             pytorch: 2.7.0
             axolotl_extras:
+          - cuda: 128
+            cuda_version: 12.8.1
+            python_version: "3.11"
+            pytorch: 2.7.0
+            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout
@@ -94,6 +99,11 @@ jobs:
             python_version: "3.11"
             pytorch: 2.7.0
             axolotl_extras:
+          - cuda: 128
+            cuda_version: 12.8.1
+            python_version: "3.11"
+            pytorch: 2.7.0
+            axolotl_extras:
     runs-on: axolotl-gpu-runner
     steps:
       - name: Checkout

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

@@ -32,21 +32,25 @@ jobs:
             pytorch: 2.6.0
             axolotl_extras: vllm
             num_gpus: 2
-            nightly_build: "true"
+          - cuda: 126
+            cuda_version: 12.6.3
+            python_version: "3.11"
+            pytorch: 2.6.0
+            axolotl_extras:
+            suffix: "-hopper"
+            num_gpus: 2
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"
             pytorch: 2.5.1
             axolotl_extras:
             num_gpus: 2
-            nightly_build: "true"
           - cuda: 126
             cuda_version: 12.6.3
             python_version: "3.11"
             pytorch: 2.7.0
             axolotl_extras:
             num_gpus: 2
-            nightly_build: "true"
     runs-on: [self-hosted, modal]
     timeout-minutes: 120
     steps:
@@ -68,7 +72,6 @@ jobs:
           echo "AXOLOTL_EXTRAS=${{ matrix.axolotl_extras}}" >> $GITHUB_ENV
           echo "CUDA=${{ matrix.cuda }}" >> $GITHUB_ENV
           echo "N_GPUS=${{ matrix.num_gpus }}" >> $GITHUB_ENV
-          echo "NIGHTLY_BUILD=${{ matrix.nightly_build }}" >> $GITHUB_ENV
           echo "CODECOV_TOKEN=${{ secrets.CODECOV_TOKEN }}" >> $GITHUB_ENV
       - name: Run tests job on Modal
         run: |

.github/workflows/tests.yml

@@ -295,6 +295,7 @@ jobs:
           find "$(pip cache dir)/http-v2" -type f -mtime +14 -exec rm {} \;
 
   docker-e2e-tests-1st:
+    # Run this job first as a gate for running the remainder of the test matrix
     if: ${{ ! contains(github.event.commits[0].message, '[skip e2e]') && github.repository_owner == 'axolotl-ai-cloud' }}
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
@@ -341,6 +342,8 @@ jobs:
     # this job needs to be run on self-hosted GPU runners...
     runs-on: [self-hosted, modal]
     timeout-minutes: 90
+    # Only run the remainder of the matrix if the first e2e check passed;
+    # this is to save on wasted compute costs for known failures that get caught in the first run
     needs: [pre-commit, pytest, docker-e2e-tests-1st]
 
     strategy:
@@ -365,6 +368,12 @@ jobs:
             pytorch: 2.7.0
             num_gpus: 1
             axolotl_extras:
+          - cuda: 128
+            cuda_version: 12.8.1
+            python_version: "3.11"
+            pytorch: 2.7.0
+            num_gpus: 1
+            axolotl_extras:
     steps:
       - name: Checkout
         uses: actions/checkout@v4

cicd/Dockerfile.jinja

@@ -32,6 +32,11 @@ RUN if [ "$NIGHTLY_BUILD" = "true" ] ; then \
     fi
 
 RUN pip install packaging==23.2 setuptools==75.8.0
+RUN if [ "$PYTORCH_VERSION" = "2.6.0" ] && [ "$CUDA" = "126" ] ; then \
+        curl -L -O https://d1dttdx32dkk5p.cloudfront.net/fa3/cu${CUDA}/torch-${PYTORCH_VERSION}/flash_attn_3-3.0.0b1-cp311-cp311-linux_x86_64.whl; \
+        pip3 install --no-cache-dir flash_attn_3-3.0.0b1-cp311-cp311-linux_x86_64.whl; \
+        rm flash_attn_3-3.0.0b1-cp311-cp311-linux_x86_64.whl; \
+    fi
 RUN if [ "$AXOLOTL_EXTRAS" != "" ] ; then \
         pip install --no-build-isolation -e .[deepspeed,flash-attn,ring-flash-attn,optimizers,ray,$AXOLOTL_EXTRAS] $AXOLOTL_ARGS; \
     else \

cicd/multigpu.py

@@ -70,7 +70,7 @@ def run_cmd(cmd: str, run_folder: str):
     image=cicd_image,
     gpu=GPU_CONFIG,
     timeout=90 * 60,
-    cpu=8.0,
+    cpu=16.0,
     memory=131072 * N_GPUS,
     volumes=VOLUME_CONFIG,
 )

docker/Dockerfile-base

@@ -1,5 +1,5 @@
-ARG CUDA_VERSION="11.8.0"
-ARG CUDNN_VERSION="8"
+ARG CUDA_VERSION="12.4.1"
+ARG CUDNN_VERSION=""
 ARG UBUNTU_VERSION="22.04"
 ARG MAX_JOBS=4
 
@@ -7,16 +7,16 @@ FROM nvidia/cuda:$CUDA_VERSION-cudnn$CUDNN_VERSION-devel-ubuntu$UBUNTU_VERSION A
 
 ENV PATH="/root/miniconda3/bin:${PATH}"
 
-ARG PYTHON_VERSION="3.10"
-ARG PYTORCH_VERSION="2.1.2"
-ARG CUDA="118"
+ARG PYTHON_VERSION="3.11"
+ARG PYTORCH_VERSION="2.5.1"
+ARG CUDA="124"
 ARG TORCH_CUDA_ARCH_LIST="7.0 7.5 8.0 8.6 9.0+PTX"
 
 ENV PYTHON_VERSION=$PYTHON_VERSION
 ENV TORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST
 
 RUN apt-get update \
-    && apt-get install -y wget git build-essential ninja-build git-lfs libaio-dev pkg-config && rm -rf /var/lib/apt/lists/* \
+    && apt-get install -y wget git build-essential ninja-build git-lfs libaio-dev pkg-config curl && rm -rf /var/lib/apt/lists/* \
     && wget \
     https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh \
     && mkdir /root/.conda \
@@ -38,6 +38,10 @@ RUN git lfs install --skip-repo && \
     # The base image ships with `pydantic==1.8.2` which is not working
     pip3 install -U --no-cache-dir pydantic==1.10.10
 
-RUN if [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \
+RUN if [ "$TORCH_CUDA_ARCH_LIST" = "9.0+PTX" ] ; then \
+        curl -L -O https://d1dttdx32dkk5p.cloudfront.net/fa3/cu${CUDA}/torch-${PYTORCH_VERSION}/flash_attn_3-3.0.0b1-cp311-cp311-linux_x86_64.whl; \
+        pip3 install --no-cache-dir flash_attn_3-3.0.0b1-cp311-cp311-linux_x86_64.whl; \
+        rm flash_attn_3-3.0.0b1-cp311-cp311-linux_x86_64.whl; \
+    elif [ "$PYTORCH_VERSION" = "2.7.0" ] ; then \
         pip3 install flash-attn==2.7.4.post1; \
     fi

docs/config.qmd

@@ -332,8 +332,6 @@ dataset_shard_idx:
 # The maximum length of an input to train with, this should typically be less than 2048
 # as most models have a token/context limit of 2048
 sequence_len: 2048
-# How to handle sequences that overflow the sequence_len: 'drop' (default, removes sample) or 'truncate' (cuts off excess tokens).
-sequence_len_overflow_handling: drop
 # Pad inputs so each step uses constant sized buffers
 # This will reduce memory fragmentation and may prevent OOMs, by re-using memory more efficiently
 pad_to_sequence_len:
@@ -635,7 +633,9 @@ weight_decay:
 # adamw hyperparams
 adam_beta1:
 adam_beta2:
+adam_beta3:  # only used for CAME Optimizer
 adam_epsilon:
+adam_epsilon2:  # only used for CAME Optimizer
 # Gradient clipping max norm
 max_grad_norm:
 

docs/getting-started.qmd

@@ -104,7 +104,7 @@ the `alpaca` dataset format, which has the following format:
 Please see our [Dataset Formats](dataset-formats) for more dataset formats and how to
 format them.
 
-2. Prepare your JSONL data in the specified format (in this case, the expected `alpaca
+2. Prepare your JSONL data in the specified format (in this case, the expected `alpaca`
 format):
 
 ```json
@@ -120,6 +120,12 @@ axolotl train my_training.yml
 
 ## Common Tasks {#sec-common-tasks}
 
+::: {.callout-tip}
+
+The same yaml file is used for training, inference, and merging.
+
+:::
+
 ### Testing Your Model {#sec-testing}
 
 After training, test your model:
@@ -128,6 +134,16 @@ After training, test your model:
 axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out"
 ```
 
+More details can be found in [Inference](inference.qmd).
+
+### Using a UI {#sec-ui}
+
+Launch a Gradio interface:
+
+```bash
+axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out" --gradio
+```
+
 ### Preprocessing Data {#sec-preprocessing}
 
 For large datasets, preprocess first:
@@ -136,14 +152,22 @@ For large datasets, preprocess first:
 axolotl preprocess my_training.yml
 ```
 
-### Using a UI {#sec-ui}
+Please make sure to set `dataset_prepared_path: ` in your config to set the path to save the prepared dataset.
 
-Launch a Gradio interface:
+More details can be found in [Dataset Preprocessing](dataset_preprocessing.qmd).
+
+### Merging LoRA weights {#sec-merging-lora}
+
+To merge the LoRA weights back into the base model, run:
 
 ```bash
-axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out" --gradio
+axolotl merge-lora my_training.yml --lora-model-dir="./outputs/lora-out"
 ```
 
+The merged model will be saved in the `{output_dir}/merged` directory.
+
+More details can be found in [Merging LoRA weights](inference.qmd#sec-merging).
+
 ## Next Steps {#sec-next-steps}
 
 Now that you have the basics, you might want to:
@@ -156,6 +180,7 @@ Now that you have the basics, you might want to:
 Check our other guides for details on these topics:
 
 - [Configuration Guide](config.qmd) - Full configuration options
+- [Dataset Loading](dataset-loading.qmd) - Loading datasets from various sources
 - [Dataset Formats](dataset-formats) - Working with different data formats
 - [Multi-GPU Training](multi-gpu.qmd)
 - [Multi-Node Training](multi-node.qmd)

src/axolotl/core/trainer_builder.py

@@ -387,8 +387,12 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
             training_arguments_kwargs["adam_beta1"] = self.cfg.adam_beta1
         if self.cfg.adam_beta2:
             training_arguments_kwargs["adam_beta2"] = self.cfg.adam_beta2
+        if self.cfg.adam_beta3:
+            training_arguments_kwargs["adam_beta3"] = self.cfg.adam_beta3
         if self.cfg.adam_epsilon:
             training_arguments_kwargs["adam_epsilon"] = self.cfg.adam_epsilon
+        if self.cfg.adam_epsilon2:
+            training_arguments_kwargs["adam_epsilon2"] = self.cfg.adam_epsilon2
         if self.cfg.max_grad_norm:
             training_arguments_kwargs["max_grad_norm"] = self.cfg.max_grad_norm
 
@@ -713,7 +717,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
 
                 beta1 = training_arguments_kwargs.get("adam_beta1", 0.9)
                 beta2 = training_arguments_kwargs.get("adam_beta2", 0.999)
-                beta3 = training_arguments_kwargs.get("adam_beta2", 0.9999)
+                beta3 = training_arguments_kwargs.get("adam_beta3", 0.9999)
                 eps1 = training_arguments_kwargs.get("adam_epsilon", 1e-30)
                 eps2 = training_arguments_kwargs.get("adam_epsilon2", 1e-16)
                 adam_kwargs["betas"] = (beta1, beta2, beta3)
@@ -1170,7 +1174,8 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
         if self.eval_dataset:
             trainer_kwargs["eval_dataset"] = self.eval_dataset
         if self.cfg.adapter and self.peft_config:
-            trainer_kwargs["peft_config"] = self.peft_config
+            if self.cfg.rl is not RLType.GRPO:
+                trainer_kwargs["peft_config"] = self.peft_config
         if self.cfg.precompute_ref_log_probs is not None:
             trainer_kwargs["precompute_ref_log_probs"] = (
                 self.cfg.precompute_ref_log_probs

src/axolotl/core/trainers/grpo/trainer.py

@@ -3,7 +3,6 @@
 # pylint: disable=too-many-lines,duplicate-code,protected-access,no-member
 
 import warnings
-from contextlib import nullcontext
 from typing import Any
 
 import datasets
@@ -14,7 +13,7 @@ from accelerate.utils import (
     broadcast_object_list,
     gather,
     gather_object,
-    is_peft_model,
+    is_peft_available,
 )
 from datasets import Dataset, IterableDataset
 from torch import nn
@@ -30,15 +29,13 @@ from transformers import (
     TrainerCallback,
 )
 from transformers.trainer_utils import seed_worker
-from transformers.utils import is_peft_available
 from trl import GRPOTrainer
 from trl.data_utils import (
     apply_chat_template,
     is_conversational,
     maybe_apply_chat_template,
 )
-from trl.extras.profiling import profiling_context, profiling_decorator
-from trl.import_utils import is_deepspeed_available
+from trl.extras.profiling import profiling_context
 from trl.models import unwrap_model_for_generation
 from trl.trainer.grpo_config import GRPOConfig
 from trl.trainer.grpo_trainer import RewardFunc, nanstd
@@ -52,62 +49,12 @@ if is_peft_available():
     # pylint: disable=unused-import
     from peft import PeftConfig
 
-if is_deepspeed_available():
-    import deepspeed
-
 
 class AxolotlGRPOTrainer(RngLoaderMixin, SchedulerMixin, GRPOTrainer):
     """Extend the base GRPOTrainer for axolotl helpers"""
 
     _tag_names = ["trl", "grpo", "axolotl"]
 
-    @profiling_decorator
-    def _move_model_to_vllm(self):
-        # For DeepSpeed ZeRO-3, we need to gather all parameters before operations
-        deepspeed_plugin = self.accelerator.state.deepspeed_plugin
-        zero_stage_3 = deepspeed_plugin is not None and deepspeed_plugin.zero_stage == 3
-        gather_if_zero3 = (
-            deepspeed.zero.GatheredParameters if zero_stage_3 else nullcontext
-        )
-
-        if is_peft_model(self.model):
-            # With PEFT and DeepSpeed ZeRO Stage 3, we must gather the full model at once before merging, as merging
-            # adapters in a sharded manner is not supported.
-            with gather_if_zero3(list(self.model.parameters())):
-                self.model.merge_adapter()
-
-                # Update vLLM weights while parameters are gathered
-                for name, param in self.model.named_parameters():
-                    # When using PEFT, we need to recover the original parameter name and discard some parameters
-                    name = (
-                        name.removeprefix("base_model.model.")
-                        .removeprefix("base_model.model.")
-                        .replace(".base_layer", "")
-                    )
-                    if self.model.prefix in name:
-                        continue
-                    # When module to save, remove its prefix and discard the original module
-                    if "original_module" in name:
-                        continue
-                    name = name.replace("modules_to_save.default.", "")
-
-                    if self.accelerator.is_main_process:
-                        self.vllm_client.update_named_param(name, param.data)
-
-                # Unmerge adapters while parameters are still gathered
-                self.model.unmerge_adapter()
-                # Parameters will automatically be repartitioned when exiting the context
-        else:
-            # For non-PEFT models, simply gather and update each parameter individually.
-            for name, param in self.model.named_parameters():
-                with gather_if_zero3([param]):
-                    if self.accelerator.is_main_process:
-                        self.vllm_client.update_named_param(name, param.data)
-
-        # Reset cache on main process
-        if self.accelerator.is_main_process:
-            self.vllm_client.reset_prefix_cache()
-
 
 class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
     """Extend the base GRPOTrainer for sequence parallelism handling"""

src/axolotl/core/training_args.py

@@ -227,6 +227,19 @@ class AxolotlTrainingMixins:
         },
     )
 
+    adam_beta3: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "The beta3 hyperparameter used in some optimizers such as CAME"
+        },
+    )
+    adam_epsilon2: Optional[float] = field(
+        default=None,
+        metadata={
+            "help": "The epsilon2 hyperparameter used in some optimizers such as CAME"
+        },
+    )
+
     # multi-modal section
 
     image_size: int | tuple[int, int] | None = field(

src/axolotl/integrations/cut_cross_entropy/monkeypatch/cohere.py

@@ -20,25 +20,15 @@ from cut_cross_entropy.transformers.utils import (
 from transformers.cache_utils import Cache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.cohere.modeling_cohere import (
-    _CONFIG_FOR_DOC,
-    COHERE_INPUTS_DOCSTRING,
     KwargsForCausalLM,
 )
 from transformers.processing_utils import Unpack
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 from transformers.utils.deprecation import deprecate_kwarg
 
 _PATCH_OPTS: PatchOptions | None = None
 
 
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(COHERE_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward(
     self,
     input_ids: torch.LongTensor | None = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma.py

@@ -17,25 +17,15 @@ from cut_cross_entropy.transformers.utils import (
 from transformers.cache_utils import Cache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.gemma.modeling_gemma import (
-    _CONFIG_FOR_DOC,
-    GEMMA_INPUTS_DOCSTRING,
     KwargsForCausalLM,
 )
 from transformers.processing_utils import Unpack
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 from transformers.utils.deprecation import deprecate_kwarg
 
 _PATCH_OPTS: PatchOptions | None = None
 
 
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward(
     self,
     input_ids: torch.LongTensor | None = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/gemma3.py

@@ -20,15 +20,11 @@ from torch import nn
 from transformers.cache_utils import Cache, HybridCache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.gemma3.modeling_gemma3 import (
-    _CONFIG_FOR_DOC,
-    GEMMA3_INPUTS_DOCSTRING,
     Gemma3CausalLMOutputWithPast,
     logger,
 )
 from transformers.utils import (
-    add_start_docstrings_to_model_forward,
     is_torchdynamo_compiling,
-    replace_return_docstrings,
 )
 from transformers.utils.deprecation import deprecate_kwarg
 
@@ -38,10 +34,6 @@ _PATCH_OPTS: PatchOptions | None = None
 
 
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(GEMMA3_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward(
     self,
     input_ids: torch.LongTensor | None = None,
@@ -170,10 +162,6 @@ def cce_forward(
 
 
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(GEMMA3_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=Gemma3CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward_multimodal(
     self,
     input_ids: torch.LongTensor | None = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama.py

@@ -19,15 +19,9 @@ from transformers.modeling_outputs import (
     CausalLMOutputWithPast,
 )
 from transformers.models.llama.modeling_llama import (
-    _CONFIG_FOR_DOC,
-    LLAMA_INPUTS_DOCSTRING,
     KwargsForCausalLM,
 )
 from transformers.processing_utils import Unpack
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 from transformers.utils.deprecation import deprecate_kwarg
 from transformers.utils.generic import can_return_tuple
 
@@ -36,10 +30,6 @@ _PATCH_OPTS: PatchOptions | None = None
 
 @can_return_tuple
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward(
     self,
     input_ids: Optional[torch.LongTensor] = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/llama4.py

@@ -16,22 +16,12 @@ from torch import nn
 from transformers.cache_utils import Cache
 from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.llama4.modeling_llama4 import (
-    _CONFIG_FOR_DOC,
-    LLAMA4_INPUTS_DOCSTRING,
     Llama4CausalLMOutputWithPast,
 )
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 
 _PATCH_OPTS: PatchOptions | None = None
 
 
-@add_start_docstrings_to_model_forward(LLAMA4_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward(
     self,
     input_ids: torch.LongTensor | None = None,
@@ -160,9 +150,6 @@ def cce_forward(
     )
 
 
-@replace_return_docstrings(
-    output_type=Llama4CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward_multimodal(
     self,
     input_ids: torch.LongTensor | None = None,  # type: ignore

src/axolotl/integrations/cut_cross_entropy/monkeypatch/mistral3.py

@@ -19,15 +19,11 @@ from transformers.models.mistral3.modeling_mistral3 import (
     Mistral3CausalLMOutputWithPast,
 )
 from transformers.models.mistral.modeling_mistral import (
-    _CONFIG_FOR_DOC,
-    MISTRAL_INPUTS_DOCSTRING,
     KwargsForCausalLM,
 )
 from transformers.processing_utils import Unpack
 from transformers.utils import (
-    add_start_docstrings_to_model_forward,
     is_torchdynamo_compiling,
-    replace_return_docstrings,
 )
 from transformers.utils.deprecation import deprecate_kwarg
 
@@ -35,10 +31,6 @@ _PATCH_OPTS: PatchOptions | None = None
 
 
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward(
     self,
     input_ids: torch.LongTensor | None = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_moe.py

@@ -13,16 +13,10 @@ from cut_cross_entropy.transformers.utils import (
     apply_lce,
 )
 from transformers.models.qwen2_moe.modeling_qwen2_moe import (
-    _CONFIG_FOR_DOC,
-    QWEN2MOE_INPUTS_DOCSTRING,
     MoeCausalLMOutputWithPast,
     MoeModelOutputWithPast,
     load_balancing_loss_func,
 )
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 from transformers.utils.deprecation import deprecate_kwarg
 from transformers.utils.generic import can_return_tuple
 
@@ -31,10 +25,6 @@ _PATCH_OPTS: PatchOptions | None = None
 
 @can_return_tuple
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def forward(
     self,
     input_ids: Optional[torch.LongTensor] = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen2_vl.py

@@ -14,22 +14,12 @@ from cut_cross_entropy.transformers.utils import (
 )
 from torch.nn import CrossEntropyLoss
 from transformers.models.qwen2_vl.modeling_qwen2_vl import (
-    _CONFIG_FOR_DOC,
-    QWEN2_VL_INPUTS_DOCSTRING,
     Qwen2VLCausalLMOutputWithPast,
 )
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 
 _PATCH_OPTS: PatchOptions | None = None
 
 
-@add_start_docstrings_to_model_forward(QWEN2_VL_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=Qwen2VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def cce_forward_multimodal(
     self,
     input_ids: Optional[torch.LongTensor] = None,

src/axolotl/integrations/cut_cross_entropy/monkeypatch/qwen3_moe.py

@@ -12,20 +12,13 @@ from cut_cross_entropy.transformers.utils import (
     TransformersModelT,
     apply_lce,
 )
-from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.models.qwen3_moe.modeling_qwen3_moe import (
-    _CONFIG_FOR_DOC,
-    QWEN3_MOE_INPUTS_DOCSTRING,
     KwargsForCausalLM,
     MoeCausalLMOutputWithPast,
     MoeModelOutputWithPast,
     load_balancing_loss_func,
 )
 from transformers.processing_utils import Unpack
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 from transformers.utils.deprecation import deprecate_kwarg
 from transformers.utils.generic import can_return_tuple
 
@@ -34,10 +27,6 @@ _PATCH_OPTS: PatchOptions | None = None
 
 @can_return_tuple
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def forward(
     self,
     input_ids: Optional[torch.LongTensor] = None,

src/axolotl/integrations/liger/models/deepseekv2.py

@@ -14,10 +14,6 @@ from torch.nn import CrossEntropyLoss
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 
-# @add_start_docstrings_to_model_forward(DeepseekV2_INPUTS_DOCSTRING)
-# @replace_return_docstrings(
-#    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-# )
 def lce_forward(
     self,
     input_ids: torch.LongTensor = None,

src/axolotl/integrations/liger/models/jamba.py

@@ -13,21 +13,11 @@ from liger_kernel.transformers.fused_linear_cross_entropy import (
 from torch.nn import CrossEntropyLoss
 from transformers.modeling_outputs import MoeCausalLMOutputWithPast
 from transformers.models.jamba.modeling_jamba import (
-    _CONFIG_FOR_DOC,
-    JAMBA_INPUTS_DOCSTRING,
     HybridMambaAttentionDynamicCache,
     load_balancing_loss_func,
 )
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 
 
-@add_start_docstrings_to_model_forward(JAMBA_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def lce_forward(
     self,
     input_ids: torch.LongTensor = None,

src/axolotl/monkeypatch/gemma3.py

@@ -7,24 +7,16 @@ from typing import Optional, Tuple, Union
 import torch
 from transformers.cache_utils import Cache
 from transformers.models.gemma3.modeling_gemma3 import (
-    _CONFIG_FOR_DOC,
-    GEMMA3_INPUTS_DOCSTRING,
     Gemma3CausalLMOutputWithPast,
     logger,
 )
 from transformers.utils import (
-    add_start_docstrings_to_model_forward,
     is_torchdynamo_compiling,
-    replace_return_docstrings,
 )
 from transformers.utils.deprecation import deprecate_kwarg
 
 
 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(GEMMA3_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=Gemma3CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
 def new_forward(
     self,
     input_ids: torch.LongTensor = None,

src/axolotl/utils/callbacks/mlflow_.py

@@ -1,6 +1,7 @@
 """MLFlow module for trainer callbacks"""
 
 import logging
+import os
 from shutil import copyfile
 from tempfile import NamedTemporaryFile
 from typing import TYPE_CHECKING
@@ -16,6 +17,11 @@ if TYPE_CHECKING:
 LOG = logging.getLogger("axolotl.callbacks")
 
 
+def should_log_artifacts() -> bool:
+    truths = ["TRUE", "1", "YES"]
+    return os.getenv("HF_MLFLOW_LOG_ARTIFACTS", "FALSE").upper() in truths
+
+
 class SaveAxolotlConfigtoMlflowCallback(TrainerCallback):
     # pylint: disable=duplicate-code
     """Callback to save axolotl config to mlflow"""
@@ -32,13 +38,18 @@ class SaveAxolotlConfigtoMlflowCallback(TrainerCallback):
     ):
         if is_main_process():
             try:
-                with NamedTemporaryFile(
-                    mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
-                ) as temp_file:
-                    copyfile(self.axolotl_config_path, temp_file.name)
-                    mlflow.log_artifact(temp_file.name, artifact_path="")
+                if should_log_artifacts():
+                    with NamedTemporaryFile(
+                        mode="w", delete=False, suffix=".yml", prefix="axolotl_config_"
+                    ) as temp_file:
+                        copyfile(self.axolotl_config_path, temp_file.name)
+                        mlflow.log_artifact(temp_file.name, artifact_path="")
+                        LOG.info(
+                            "The Axolotl config has been saved to the MLflow artifacts."
+                        )
+                else:
                     LOG.info(
-                        "The Axolotl config has been saved to the MLflow artifacts."
+                        "Skipping logging artifacts to MLflow (hf_mlflow_log_artifacts is false)"
                     )
             except (FileNotFoundError, ConnectionError) as err:
                 LOG.warning(f"Error while saving Axolotl config to MLflow: {err}")

src/axolotl/utils/data/pretraining.py

@@ -11,7 +11,6 @@ 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.samplers import MultipackBatchSampler, get_dataset_lengths
 from axolotl.utils.trainer import process_pretraining_datasets_for_packing
 
@@ -251,22 +250,6 @@ def encode_packed_pretraining(
     # pylint: disable=duplicate-code
     # tokenize all the examples
     # rows get split with stride (overlap)
-    """
-    Encodes and packs input examples into fixed-length batches for pretraining with optional multipack attention.
-    
-    Wraps and processes input examples into a dataset, applies sequence packing with configurable overflow handling, and batches the data using a multipack sampler. Each batch is collated and features are aggregated into lists keyed by feature name.
-    
-    Args:
-        collate_fn: Function to collate individual feature dictionaries into batch tensors.
-        ds_wrapper: Callable that wraps a Hugging Face Dataset for further processing.
-        examples: Dictionary of input examples to encode and pack.
-        max_seq_length: Maximum sequence length for each packed sequence.
-        batch_size: Number of sequences to pack per batch.
-        multipack_attn: If True, enables multipack attention and drops attention masks.
-    
-    Returns:
-        Dictionary where each key is a feature name and each value is a list of packed feature tensors.
-    """
     train_dataset = ds_wrapper(Dataset.from_dict(examples))[0]
 
     train_dataset = process_pretraining_datasets_for_packing(
@@ -276,10 +259,6 @@ 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", DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING
-        ),
     )
 
     sampler = MultipackBatchSampler(

src/axolotl/utils/data/rl.py

@@ -72,6 +72,7 @@ def map_dataset(cfg, data_set, ds_transform_fn, tokenizer, **map_kwargs):
     data_set = data_set.map(
         ds_transform_fn,
         desc="Mapping RL Dataset",
+        num_proc=cfg.dataset_processes,
         **map_kwargs,
     )
 
@@ -79,33 +80,8 @@ def map_dataset(cfg, data_set, ds_transform_fn, tokenizer, **map_kwargs):
 
 
 def drop_long_rl_seq(
-    sample,
-    rl,
-    tokenizer,
-    sequence_len,
-    handling="drop",  # Use the default handling mode
+    sample, rl, tokenizer, sequence_len  # pylint: disable=invalid-name
 ):
-    """
-    Handles samples exceeding a maximum sequence length for various RL dataset types by either truncating or dropping them.
-    
-    Depending on the RL type and the `handling` mode, this function either truncates response fields to fit within the specified sequence length or determines whether the sample should be dropped. For DPO, IPO, ORPO, and SIMPO types, both "chosen" and "rejected" responses are considered; for KTO, the "completion" is considered. For GRPO, samples are always retained. If truncation is not possible (e.g., the prompt alone exceeds the limit), the sample is returned unchanged for mapping, or dropped during filtering.
-    
-    Args:
-        sample: A dictionary representing a single dataset sample.
-        rl: The RLType indicating the dataset type.
-        tokenizer: The tokenizer used to compute token lengths and perform truncation.
-        sequence_len: The maximum allowed sequence length.
-        handling: Specifies how to handle overlong sequences ("drop" or "truncate").
-    
-    Returns:
-        For "truncate": The modified sample with responses truncated as needed, or the original sample if truncation is not possible.
-        For "drop": True if the sample fits within the sequence length, otherwise False.
-    
-    Raises:
-        ValueError: If required keys are missing for the specified RL type, or if the RL type is unknown.
-    """
-    result = None
-
     if rl in (RLType.DPO, RLType.IPO, RLType.ORPO, RLType.SIMPO):
         if not (
             sample.get("prompt") and sample.get("chosen") and sample.get("rejected")
@@ -122,65 +98,11 @@ def drop_long_rl_seq(
         len_chosen = len(tokenizer(chosen, add_special_tokens=False)["input_ids"])
         len_rejected = len(tokenizer(rejected, add_special_tokens=False)["input_ids"])
 
-        # Truncate first, then drop if still invalid (although truncate should handle it)
-        if handling == "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
+        return (len_prompt + len_chosen) <= sequence_len and (
+            len_prompt + len_rejected
+        ) <= sequence_len
 
-                if max_response_len <= 0:
-                    # Prompt is already too long, behavior depends on handling
-                    # If truncate is chosen, we technically can't truncate, but drop seems harsh.
-                    # Returning the sample might be unexpected. Let's stick to the filter logic
-                    # which would drop this in the `filter` step later if needed.
-                    # For now, return sample to map, or False to filter.
-                    # Let's simplify: truncate *should* result in a valid sample if possible.
-                    # If prompt >= seq_len, truncate won't work. Filter will catch this later.
-                    # So, if max_response_len <= 0, we pass it through for map, drop for filter.
-                    # However, the filter/map logic is applied *after* this function.
-                    # This function needs to return the *modified* sample for map, or bool for filter.
-
-                    # Re-think: If handling==truncate, return the modified sample if possible.
-                    # If prompt >= seq_len, modification is impossible. What should map return?
-                    # Maybe return the original sample? But map expects *modified* sample.
-                    # Let's stick to the original logic: if prompt is too long, return False for filter
-                    # and original sample for map.
-
-                    result = (
-                        sample  # For map, let downstream handle it if still invalid?
-                    )
-                    # Or maybe return None/empty dict? Let's return sample for now.
-                    # If handling was drop, filter would remove this.
-
-                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 rl is RLType.KTO:
         if not (sample.get("prompt") and sample.get("completion")):
             raise ValueError("Prompt and completion keys are required for KTO datasets")
 
@@ -192,54 +114,15 @@ def drop_long_rl_seq(
             tokenizer(completion, add_special_tokens=False)["input_ids"]
         )
 
-        # Truncate first
-        if handling == "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
+        return (len_prompt + len_completion) <= sequence_len
 
-                if max_completion_len <= 0:
-                    # Prompt too long, return sample for map
-                    result = sample
-                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
+    if rl is RLType.GRPO:
+        return True
 
-    elif rl == RLType.GRPO:
-        # GRPO doesn't involve sequence length checks in the same way?
-        # The original code returned True for drop. What should it return for truncate?
-        # Let's assume for now it always passes.
-        result = sample if handling == "truncate" else True
-    else:
-        raise ValueError("Unknown RL type")
-
-    return result
+    raise ValueError("Unknown RL type")
 
 
 def load_prepare_preference_datasets(cfg):
-    """
-    Loads, preprocesses, and prepares preference datasets for RL training and evaluation.
-    
-    This function orchestrates the loading, transformation, sequence length handling, optional deduplication, and caching of datasets for Direct Preference Optimization (DPO) and related RL types. It supports configurable handling of overlong sequences (dropping or truncating), applies dataset-specific transformations, and manages train/validation/test splits as needed.
-    
-    Args:
-        cfg: Configuration object specifying dataset sources, RL type, tokenizer, sequence length, and processing options.
-    
-    Returns:
-        A tuple containing the prepared training and evaluation datasets.
-    """
     def load_split(dataset_cfgs, _cfg):
         split_datasets: List[Any] = []
         use_auth_token = _cfg.hf_use_auth_token
@@ -283,43 +166,25 @@ def load_prepare_preference_datasets(cfg):
                 split_datasets[i] = data_set
 
             if not cfg.skip_prepare_dataset:
-                # Determine handling mode
-                handling = cfg.get("sequence_len_overflow_handling", "drop")
-
                 drop_long = partial(
                     drop_long_rl_seq,
                     rl=_cfg.rl,
                     tokenizer=tokenizer,
                     sequence_len=cfg.sequence_len,
-                    handling=handling,  # Pass the handling mode
                 )
 
                 prior_len = len(split_datasets[i])
-
-                # Use map for truncate mode and filter for drop mode
-                if handling == "truncate":
-                    split_datasets[i] = split_datasets[i].map(
-                        drop_long,  # Function now returns modified sample or original
-                        num_proc=cfg.dataset_processes,
-                        load_from_cache_file=not cfg.is_preprocess,
-                        desc="Truncating Long Sequences",
+                split_datasets[i] = split_datasets[i].filter(
+                    drop_long,
+                    num_proc=cfg.dataset_processes,
+                    load_from_cache_file=not cfg.is_preprocess,
+                    desc="Dropping Long Sequences",
+                )
+                dropped = prior_len - len(split_datasets[i])
+                if dropped:
+                    LOG.warning(
+                        f"Dropped {dropped} long samples from dataset index {i}"
                     )
-                    # Note: Length might not change if truncation always occurs
-                    LOG.info(
-                        f"Processed dataset index {i} with truncation handling for sequence length {cfg.sequence_len}"
-                    )
-                else:  # handling == "drop"
-                    split_datasets[i] = split_datasets[i].filter(
-                        drop_long,  # Function now returns boolean
-                        num_proc=cfg.dataset_processes,
-                        load_from_cache_file=not cfg.is_preprocess,
-                        desc="Dropping Long Sequences",
-                    )
-                    dropped = prior_len - len(split_datasets[i])
-                    if dropped:
-                        LOG.warning(
-                            f"Dropped {dropped} long samples from dataset index {i}"
-                        )
 
         combined_datasets = concatenate_datasets(split_datasets)
         combined_datasets = combined_datasets.shuffle(seed=cfg.seed or 42)

src/axolotl/utils/data/sft.py

@@ -484,7 +484,7 @@ def get_dataset_wrapper(
     }
 
     LOG.info(
-        f"Loading dataset with base_type: {d_base_type} and prompt_style: {d_prompt_style}"
+        f"Loading dataset: {config_dataset['path']} with base_type: {d_base_type} and prompt_style: {d_prompt_style}"
     )
 
     if (

src/axolotl/utils/data/utils.py

@@ -13,12 +13,10 @@ from datasets import Dataset, IterableDataset
 
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.samplers.utils import get_dataset_lengths
-from axolotl.utils.trainer import truncate_or_drop_long_seq
+from axolotl.utils.trainer import drop_long_seq
 
 LOG = logging.getLogger(__name__)
 
-DEFAULT_SEQUENCE_LEN_OVERFLOW_HANDLING = "drop"
-
 
 class RetryStrategy(Enum):
     """
@@ -161,33 +159,16 @@ def deduplicate_and_log_datasets(
 
 
 def drop_long_seq_in_dataset(dataset: Dataset, cfg: DictDefault):
-    """
-    Processes a dataset to handle sequences exceeding a configured maximum length by either truncating or dropping them.
-    
-    If the dataset lacks an "input_ids" column, the function returns the dataset unchanged. The handling mode is determined by the configuration parameter "sequence_len_overflow_handling", defaulting to "drop". In "truncate" mode, sequences longer than the maximum length are truncated; in "drop" mode, such sequences are removed from the dataset. The function logs information about sequence lengths and the number of samples affected when applicable.
-    
-    Args:
-        dataset: The Huggingface Dataset to process.
-        cfg: Configuration object specifying sequence length parameters and handling mode.
-    
-    Returns:
-        The processed dataset with long sequences either truncated or dropped according to the configuration.
-    """
     if "input_ids" not in dataset.column_names:
         LOG.warning(
             "Dataset does not contain 'input_ids' column. Skip drop long seq. This is expected for RewardModeling."
         )
         return dataset
 
-    # Get the handling method from config, default to "drop" for backward compatibility
-    handling = cfg.get("sequence_len_overflow_handling", "drop")
-
-    # Use the new function with the specified handling mode
-    seq_handler = functools.partial(
-        truncate_or_drop_long_seq,
+    drop_long = functools.partial(
+        drop_long_seq,
         sequence_len=cfg.sequence_len,
         min_sequence_len=cfg.min_sample_len,
-        handling=handling,
     )
 
     try:
@@ -212,31 +193,17 @@ def drop_long_seq_in_dataset(dataset: Dataset, cfg: DictDefault):
 
     drop_long_kwargs = {}
     if filter_map_kwargs:
-        if handling == "truncate":
-            drop_long_kwargs["desc"] = "Truncating Long Sequences"
-        else:  # handling == "drop"
-            drop_long_kwargs["desc"] = "Dropping Long Sequences"
+        drop_long_kwargs["desc"] = "Dropping Long Sequences"
 
-    if handling == "truncate":
-        # Use map for truncate mode
-        dataset = dataset.map(
-            seq_handler,
-            batched=True,
-            **filter_map_kwargs,
-            **drop_long_kwargs,
-        )
-        LOG.info(f"Truncated long samples in dataset to {cfg.sequence_len} tokens")
-    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")
+    dataset = dataset.filter(
+        drop_long,
+        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/models.py

@@ -629,6 +629,49 @@ class ModelLoader:
             )
 
         if self.cfg.flash_attention:
+            use_fa3 = False
+            if self.cfg.use_flash_attention_3 is True:
+                use_fa3 = True
+            elif self.cfg.use_flash_attention_3 == "auto":
+                if torch.cuda.get_device_capability() >= (9, 0):
+                    # FA3 is only available on Hopper GPUs and newer
+                    use_fa3 = True
+                if not importlib.util.find_spec("flash_attn_interface"):
+                    use_fa3 = False
+            if use_fa3 and not importlib.util.find_spec("flash_attn_interface"):
+                # this can happen when use_flash_attention_3 is explicity set to True
+                # and flash_attn_interface is not installed
+                raise ModuleNotFoundError(
+                    "Please install the flash_attn_interface library to use Flash Attention 3.x"
+                )
+            if use_fa3 and importlib.util.find_spec("flash_attn_interface") is not None:
+                from flash_attn_interface import flash_attn_func as flash_attn_func_v3
+                from flash_attn_interface import (
+                    flash_attn_varlen_func as flash_attn_varlen_func_v3,
+                )
+
+                def flash_attn_func_v3_wrapper(*args, **kwargs):
+                    kwargs.pop("dropout_p", None)
+                    if "softmax_scale" in kwargs and len(args) >= 4:
+                        # if softmax_scale is provided, then the 3rd position is dropout_p that we need to drop
+                        args = (*args[:3],) + args[4:]
+                    return flash_attn_func_v3(*args, **kwargs)[0]
+
+                def flash_attn_varlen_func_v3_wrapper(*args, **kwargs):
+                    kwargs.pop("dropout_p", None)
+                    if "softmax_scale" in kwargs and len(args) >= 4:
+                        # if softmax_scale is provided, then the 3rd position is dropout_p that we need to drop
+                        args = (*args[:3],) + args[4:]
+                    return flash_attn_varlen_func_v3(*args, **kwargs)[0]
+
+                transformers.modeling_flash_attention_utils.flash_attn_func = (
+                    flash_attn_func_v3_wrapper
+                )
+                transformers.modeling_flash_attention_utils.flash_attn_varlen_func = (
+                    flash_attn_varlen_func_v3_wrapper
+                )
+                LOG.info("Switched to Flash Attention v3")
+
             self.patch_attention()
 
         if self.cfg.sample_packing and self.cfg.s2_attention:
@@ -699,6 +742,7 @@ class ModelLoader:
 
                 patch_mllama()
 
+            # TODO deprecate soon
             if self.model_config.model_type == "btlm":
                 from axolotl.monkeypatch.btlm_attn_hijack_flash import (
                     replace_btlm_attn_with_flash_attn,
@@ -706,6 +750,7 @@ class ModelLoader:
 
                 replace_btlm_attn_with_flash_attn(self.cfg.base_model)
 
+            # TODO deprecate soon
             if (
                 self.model_config.model_type == "stablelm_epoch"
                 and self.cfg.sample_packing

src/axolotl/utils/schemas/config.py

@@ -186,12 +186,6 @@ class AxolotlInputConfig(
     unfrozen_parameters: list[str] | None = None
 
     sequence_len: int = Field(default=512)
-    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)."
-        },
-    )
     min_sample_len: int | None = None
     max_prompt_len: int = Field(
         default=512,
@@ -239,6 +233,7 @@ class AxolotlInputConfig(
     flash_attn_fuse_qkv: bool | None = None
     flash_attn_fuse_mlp: bool | None = None
     flash_optimum: bool | None = None
+    use_flash_attention_3: Literal["auto"] | bool | None = None
 
     eager_attention: bool | None = None
 

src/axolotl/utils/trainer.py

@@ -207,18 +207,10 @@ def add_length(sample):
 
 def drop_long_seq(sample, sequence_len=2048, min_sequence_len=2):
     """
-    Determines whether samples should be kept based on sequence length constraints.
-    
-    For a single example or a batch, returns True (or a list of booleans) if each sequence's length is within the specified range; otherwise, returns False (or a list with False for out-of-range sequences).
-    
-    Args:
-        sample: A dictionary containing "input_ids" as a list of ints or a list of lists of ints.
-        sequence_len: Maximum allowed sequence length (inclusive).
-        min_sequence_len: Minimum allowed sequence length (inclusive).
-    
-    Returns:
-        True if the single example is within the length range, False otherwise.
-        For batched input, returns a list of booleans indicating which sequences are within the range.
+    Drop samples whose sequence length is either too long (> sequence_len)
+    or too short (< min_sequence_len).
+
+    Works for both single-example (list[int]) or batched (list[list[int]]).
     """
     min_sequence_len = min_sequence_len or 2
 
@@ -243,121 +235,7 @@ 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"
-):
-    """
-    Drops or truncates samples based on sequence length constraints.
-    
-    If handling is "drop", returns a boolean or list of booleans indicating whether each sample's sequence length is within the specified range. If handling is "truncate", returns the sample with sequences longer than sequence_len truncated and sequences shorter than min_sequence_len omitted. Supports both single-example and batched inputs.
-    
-    Args:
-        sample: A dictionary containing at least an "input_ids" field, representing either a single sequence or a batch of sequences.
-        sequence_len: Maximum allowed sequence length.
-        min_sequence_len: Minimum allowed sequence length.
-        handling: "drop" to filter out samples outside the range, "truncate" to truncate long sequences.
-    
-    Returns:
-        In "drop" mode, a boolean or list of booleans indicating which samples to keep. In "truncate" mode, the modified sample with sequences truncated as needed.
-    """
-    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):
-    """
-    Prepares training and evaluation datasets for sample packing and model-specific requirements.
-    
-    Removes unnecessary columns based on model type, filters out samples with no trainable tokens, and optionally adds length or position ID columns for sample packing or PoSE techniques. Returns the processed training and evaluation datasets.
-    """
     drop_attn_mask = cfg.model_config_type in ["mamba", "gemma3"]
     if drop_attn_mask:
         LOG.info("dropping attention_mask column")
@@ -492,48 +370,15 @@ 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,
-    handling="drop",
+    train_dataset, sequence_len, skip_position_ids=True, drop_attention_mask=False
 ):
-    # Define the function to use for handling sequences based on the mode
-    """
-    Processes a pretraining dataset by truncating or dropping sequences based on length.
-    
-    Depending on the handling mode, sequences longer than `sequence_len` are either truncated or dropped, and sequences shorter than `min_sequence_len` are dropped. Optionally adds position IDs and removes the attention mask column.
-    
-    Args:
-        train_dataset: The dataset to process.
-        sequence_len: Maximum allowed sequence length.
-        skip_position_ids: If False, adds position IDs to each sample.
-        drop_attention_mask: If True, removes the attention mask column.
-        handling: "drop" to remove long sequences, "truncate" to truncate them.
-    
-    Returns:
-        The processed dataset with sequences handled according to the specified mode.
-    """
-    seq_handler_fn = partial(
-        truncate_or_drop_long_seq,
-        sequence_len=sequence_len,
-        handling=handling,  # Pass handling mode
+    drop_long = partial(drop_long_seq, sequence_len=sequence_len)
+
+    train_dataset = train_dataset.filter(
+        drop_long,
+        desc="Dropping Long Sequences",
+        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/conftest.py

@@ -421,6 +421,7 @@ def temp_dir():
 
 @pytest.fixture(scope="function", autouse=True)
 def cleanup_monkeypatches():
+    import transformers.modeling_flash_attention_utils
     from transformers import Trainer
     from transformers.models.llama.modeling_llama import (  # LlamaFlashAttention2,
         LlamaAttention,
@@ -434,6 +435,19 @@ def cleanup_monkeypatches():
         Trainer._inner_training_loop  # pylint: disable=protected-access
     )
     original_trainer_training_step = Trainer.training_step
+    original_fa_func = None
+    original_fa_varlen_func = None
+    if (
+        importlib.util.find_spec("flash_attn")
+        and hasattr(transformers.modeling_flash_attention_utils, "flash_attn_func")
+        and hasattr(
+            transformers.modeling_flash_attention_utils, "flash_attn_varlen_func"
+        )
+    ):
+        original_fa_func = transformers.modeling_flash_attention_utils.flash_attn_func
+        original_fa_varlen_func = (
+            transformers.modeling_flash_attention_utils.flash_attn_varlen_func
+        )
     # monkey patches can happen inside the tests
     yield
     # Reset LlamaFlashAttention2 forward
@@ -444,6 +458,11 @@ def cleanup_monkeypatches():
         original_trainer_inner_training_loop
     )
     Trainer.training_step = original_trainer_training_step
+    if original_fa_func:
+        transformers.modeling_flash_attention_utils.flash_attn_func = original_fa_func
+        transformers.modeling_flash_attention_utils.flash_attn_varlen_func = (
+            original_fa_varlen_func
+        )
 
     # Reset other known monkeypatches
     modules_to_reset: list[tuple[str, list[str]]] = [
@@ -458,6 +477,7 @@ def cleanup_monkeypatches():
         ("transformers.trainer",),
         ("transformers", ["Trainer"]),
         ("transformers.loss.loss_utils",),
+        ("transformers.modeling_flash_attention_utils",),
     ]
     for module_name_tuple in modules_to_reset:
         module_name = module_name_tuple[0]

tests/e2e/multigpu/solo/test_grpo.py

@@ -166,7 +166,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
 """
             )
 
-    @pytest.mark.skip(reason="flaky test")
     @pytest.mark.parametrize(
         "num_gpus",
         [1, 2],
@@ -231,8 +230,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
             "NCCL_P2P_LEVEL": "LOC",
             **current_env,
             "CUDA_VISIBLE_DEVICES": "1",
-            "VLLM_DISABLE_COMPILE_CACHE": "1",
-            # "VLLM_USE_V1": "0",
         }
         vllm_process = start_vllm(
             cfg.base_model,
@@ -266,7 +263,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
         finally:
             recursive_kill(vllm_process)
 
-    @pytest.mark.skip(reason="flaky test")
     @pytest.mark.parametrize(
         "num_gpus",
         [1, 2],
@@ -325,8 +321,6 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
             "NCCL_P2P_LEVEL": "LOC",  # nccl can be brittle, assume P2P isn't reliable
             **current_env,
             "CUDA_VISIBLE_DEVICES": "1",
-            "VLLM_DISABLE_COMPILE_CACHE": "1",
-            # "VLLM_USE_V1": "0",
         }
         vllm_process = start_vllm(
             cfg.base_model,

tests/e2e/multigpu/test_llama.py

@@ -101,7 +101,13 @@ class TestMultiGPULlama:
         "gradient_accumulation_steps",
         [1, 2],
     )
-    def test_lora_ddp_packed(self, temp_dir, gradient_accumulation_steps):
+    @pytest.mark.parametrize(
+        "use_flash_attention_3",
+        [False, "auto"],
+    )
+    def test_lora_ddp_packed(
+        self, temp_dir, gradient_accumulation_steps, use_flash_attention_3
+    ):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
@@ -138,6 +144,7 @@ class TestMultiGPULlama:
                 "flash_attention": True,
                 "use_tensorboard": True,
                 "bf16": True,
+                "use_flash_attention_3": use_flash_attention_3,
             }
         )
 

tests/e2e/test_packing_loss.py

@@ -4,7 +4,6 @@ E2E tests for packed training
 
 import logging
 import os
-import unittest
 
 from transformers.utils import is_torch_bf16_gpu_available
 
@@ -14,18 +13,17 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_tensorboard, with_temp_dir
+from .utils import check_tensorboard
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 
-class TestPackedLlama(unittest.TestCase):
+class TestPackedLlama:
     """
     Test case for Packed training of llama models
     """
 
-    @with_temp_dir
     def test_loss_packed(self, temp_dir):
         # pylint: disable=duplicate-code
         cfg = DictDefault(

tests/test_data.py

@@ -3,12 +3,10 @@ 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
 
@@ -60,328 +58,11 @@ class TestEncodePretraining(unittest.TestCase):
         self.assertEqual(result["input_ids"][0][14], self.tokenizer.pad_token_id)
 
     def test_md5(self):
-        """
-        Tests that the md5 function returns the correct hash for a given string and encoding.
-        """
         self.assertEqual(md5("hello world"), "5eb63bbbe01eeed093cb22bb8f5acdc3")
         self.assertEqual(
             md5("hello world", "utf-8"), "5eb63bbbe01eeed093cb22bb8f5acdc3"
         )
 
 
-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
-        """
-        Sets up a mock tokenizer and sequence length for RL sequence length tests.
-        
-        The mock tokenizer simulates tokenization by returning input IDs equal to the input string's length and decodes tokens as repeated "x" characters. The sequence length limit is set to 20.
-        """
-        self.tokenizer = MagicMock()
-
-        def side_effect_func(
-            text, add_special_tokens=False
-        ):  # pylint: disable=unused-argument
-            """
-            Simulates tokenization by returning input IDs as a sequence of integers equal to the input text length.
-            
-            Args:
-                text: The input string to tokenize.
-                add_special_tokens: Ignored parameter included for interface compatibility.
-            
-            Returns:
-                A dictionary with 'input_ids' as a list of integers from 0 to len(text) - 1.
-            """
-            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):
-        """
-        Tests that drop_long_rl_seq returns True in drop mode for a DPO sample within the sequence length limit.
-        """
-        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):
-        """
-        Tests that in DPO drop mode, a sample is rejected when the prompt and chosen lengths exceed the sequence limit.
-        """
-        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):
-        """
-        Tests that in DPO drop mode, a sample is rejected when the prompt plus rejected response exceeds the sequence length limit.
-        """
-        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):
-        """
-        Verifies that in DPO truncate mode, samples within the sequence length limit are returned unchanged.
-        """
-        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):
-        """
-        Tests that in DPO truncate mode, if the prompt exceeds the sequence length limit,
-        the original sample is returned unchanged.
-        """
-        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):
-        """
-        Tests that in DPO truncate mode, only the 'chosen' field is truncated when it exceeds the allowed sequence length, while 'prompt' and 'rejected' remain unchanged.
-        """
-        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):
-        """
-        Tests that in DPO truncate mode, only the 'rejected' field is truncated when it exceeds the sequence length limit, while 'prompt' and 'chosen' remain unchanged.
-        """
-        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):
-        """
-        Tests that in DPO truncate mode, both 'chosen' and 'rejected' fields are truncated when their combined lengths with the prompt exceed the sequence limit.
-        
-        Verifies that both fields are truncated to fit within the allowed response length and replaced with decoded placeholder content.
-        """
-        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):
-        """
-        Tests DPO truncate mode where only the overlong response is truncated.
-        
-        Verifies that when the prompt plus one response exceeds the sequence length, only the response exceeding the maximum allowed length is truncated, while the other remains unchanged.
-        """
-        # 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):
-        """
-        Tests that drop_long_rl_seq returns True for a KTO sample within the sequence length limit.
-        """
-        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):
-        """
-        Tests that drop_long_rl_seq returns False when a KTO sample exceeds the sequence length limit in drop mode.
-        """
-        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):
-        """
-        Tests that KTO truncate mode returns the original sample unchanged when the combined prompt and completion length does not exceed the sequence limit.
-        """
-        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):
-        """
-        Tests that in KTO truncate mode, if the prompt exceeds the sequence length limit, the original sample is returned unchanged.
-        """
-        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):
-        """
-        Tests that in KTO truncate mode, the completion is truncated when the combined prompt and completion exceed the sequence length limit.
-        
-        Verifies that the prompt remains unchanged and the completion is truncated to fit within the allowed length, with the truncated completion replaced by decoded "x" characters.
-        """
-        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):
-        """
-        Tests that a ValueError is raised when required keys are missing for DPO samples.
-        
-        Verifies that the function raises an error if the sample does not contain 'chosen' and 'rejected' keys.
-        """
-        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):
-        """
-        Tests that a ValueError is raised when required keys are missing for RL type "kto".
-        
-        Verifies that calling drop_long_rl_seq with a sample missing the "completion" key raises
-        a ValueError with the expected error message.
-        """
-        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):
-        """
-        Tests that a ValueError is raised when an unknown RL type is provided to drop_long_rl_seq.
-        """
-        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):
-        """
-        Tests that drop_long_rl_seq in GRPO drop mode always returns True, regardless of input.
-        """
-        sample = {}
-        result = drop_long_rl_seq(
-            sample, "grpo", self.tokenizer, self.sequence_len, handling="drop"
-        )
-        self.assertTrue(result)
-
-    def test_grpo_truncate(self):
-        """
-        Tests that in truncate mode for RL type "grpo", the original sample is returned unchanged.
-        """
-        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

@@ -1,175 +0,0 @@
-"""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
-        """
-        Sets up default sequence length parameters for the test cases.
-        """
-        self.sequence_len = 10
-        self.min_sequence_len = 3
-
-    def test_drop_mode_single(self):
-        """
-        Verifies that 'drop' mode correctly filters single sequence examples based on length.
-        
-        Tests that sequences shorter than the minimum, longer than the maximum, or empty are dropped,
-        while sequences within the valid length range are kept.
-        """
-        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):
-        """
-        Tests that 'truncate_or_drop_long_seq' correctly truncates or preserves single examples in "truncate" mode.
-        
-        Verifies that sequences longer than the maximum length are truncated, while sequences that are too short, empty, or within the valid range remain unchanged.
-        """
-        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):
-        """
-        Tests that the "drop" handling mode correctly filters batched input sequences based on length constraints.
-        
-        Verifies that sequences shorter than the minimum length, longer than the maximum length, or empty are dropped (returns False), while sequences within the valid range are kept (returns True).
-        """
-        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):
-        """
-        Tests that batched examples are correctly truncated in "truncate" mode.
-        
-        Verifies that sequences in both "input_ids" and "labels" longer than the maximum
-        allowed length are truncated, while sequences that are too short or empty remain
-        unchanged.
-        """
-        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()