avoid crash/oom on train end

.github/workflows/main.yml

@@ -31,11 +31,6 @@ 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
@@ -99,11 +94,6 @@ 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/tests.yml

@@ -295,7 +295,6 @@ 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]
@@ -342,8 +341,6 @@ 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:
@@ -368,12 +365,6 @@ 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

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,12 +120,6 @@ 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:
@@ -134,16 +128,6 @@ 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:
@@ -152,22 +136,14 @@ For large datasets, preprocess first:
 axolotl preprocess my_training.yml
 ```
 
-Please make sure to set `dataset_prepared_path: ` in your config to set the path to save the prepared dataset.
+### Using a UI {#sec-ui}
 
-More details can be found in [Dataset Preprocessing](dataset_preprocessing.qmd).
+Launch a Gradio interface:
-
-### Merging LoRA weights {#sec-merging-lora}
-
-To merge the LoRA weights back into the base model, run:
 
 ```bash
-axolotl merge-lora my_training.yml --lora-model-dir="./outputs/lora-out"
+axolotl inference my_training.yml --lora-model-dir="./outputs/lora-out" --gradio
 ```
 
-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:
@@ -180,7 +156,6 @@ 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/trainers/base.py

@@ -156,9 +156,6 @@ class AxolotlTrainer(
         Helper method to get the sampler for evaluation. Handles sequence parallelism
         and sample packing cases.
 
-        Args:
-            eval_dataset: Evaluation dataset.
-
         Returns:
             If the dataset is non-empty, a sampler is returned, the type of which
                 depends on the passed training args.
@@ -240,6 +237,9 @@ class AxolotlTrainer(
             self.accelerator.even_batches = False
 
         # Return unprepared dataloader if using sequence parallelism
+        # TODO(djsaunde): We might be able to use `accelerate`'s dataloader preparation
+        # if we use `dispatch_batches` and `slice_fn_for_dispatch` properly (i.e.,
+        # slice each batch along the sequence dimension).
         if self.args.sequence_parallel_degree > 1:
             return dataloader
 

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

@@ -1,25 +1,33 @@
-"""DPO trainer for Axolotl"""
+"""
+DPO trainer for axolotl
+"""
 
 import gc
+import random
 from functools import wraps
-from typing import Any, Dict, Union
+from typing import Any, Dict, Optional, Union
 
+import pandas as pd
 import torch
-from datasets import Dataset
+import wandb
+from accelerate import PartialState
+from datasets import Dataset, IterableDataset
 from peft.optimizers import create_loraplus_optimizer
 from torch import nn
-from torch.utils.data import Sampler
+from torch.utils.data import DataLoader
 from transformers import (
+    BaseImageProcessor,
+    FeatureExtractionMixin,
+    PreTrainedTokenizerBase,
+    ProcessorMixin,
     Trainer,
 )
+from transformers.trainer_utils import EvalLoopOutput
 from transformers.utils import is_sagemaker_mp_enabled
-from trl import DPOTrainer
+from trl import DPOConfig, DPOTrainer, maybe_apply_chat_template, maybe_extract_prompt
+from trl.trainer.utils import log_table_to_comet_experiment
 
-from axolotl.core.trainers.mixins import (
+from axolotl.core.trainers.mixins import RngLoaderMixin, SchedulerMixin
-    RngLoaderMixin,
-    SchedulerMixin,
-    SequenceParallelMixin,
-)
 from axolotl.core.trainers.utils import (
     sanitize_kwargs_for_ds_tagging,
     sanitize_kwargs_for_tagging,
@@ -29,10 +37,10 @@ if is_sagemaker_mp_enabled():
     import smdistributed.modelparallel.torch as smp
 
 
-class AxolotlDPOTrainer(
+class AxolotlDPOTrainer(RngLoaderMixin, SchedulerMixin, DPOTrainer):
-    RngLoaderMixin, SchedulerMixin, SequenceParallelMixin, DPOTrainer
+    """
-):
+    Extend the base DPOTrainer for axolotl helpers
-    """Extend the base DPOTrainer for axolotl helpers"""
+    """
 
     tag_names = ["axolotl", "dpo"]
 
@@ -87,6 +95,64 @@ class AxolotlDPOTrainer(
 
         return super().push_to_hub(*args, **kwargs)
 
+    # TODO: remove this once https://github.com/huggingface/trl/pull/3377 is in a release
+    def _prepare_dataset(
+        self,
+        dataset: Union[Dataset, IterableDataset],
+        processing_class: Union[
+            PreTrainedTokenizerBase,
+            BaseImageProcessor,
+            FeatureExtractionMixin,
+            ProcessorMixin,
+        ],
+        args: DPOConfig,
+        dataset_name: str,
+    ) -> Union[Dataset, IterableDataset]:
+        # Build the kwargs for the `map` function
+        map_kwargs: Dict[str, Any] = {"writer_batch_size": 10}
+        if isinstance(dataset, Dataset):  # IterableDataset does not support num_proc
+            map_kwargs["num_proc"] = args.dataset_num_proc
+
+        with PartialState().main_process_first():
+            # Extract prompt if needed
+            if isinstance(
+                dataset, Dataset
+            ):  # `IterableDataset.map` does not support `desc`
+                map_kwargs["desc"] = f"Extracting prompt in {dataset_name} dataset"
+            dataset = dataset.map(maybe_extract_prompt, **map_kwargs)
+
+            # Apply the chat template if needed
+            if isinstance(
+                dataset, Dataset
+            ):  # `IterableDataset.map` does not support `desc`
+                map_kwargs["desc"] = f"Applying chat template to {dataset_name} dataset"
+            dataset = dataset.map(
+                maybe_apply_chat_template,
+                fn_kwargs={"tokenizer": processing_class, "tools": args.tools},
+                **map_kwargs,
+            )
+
+            # Tokenize the dataset
+            if isinstance(
+                dataset, Dataset
+            ):  # `IterableDataset.map` does not support `desc`
+                map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"
+
+            dataset = dataset.map(
+                self.tokenize_row if not self.is_vision_model else self.process_row,
+                remove_columns=["chosen", "rejected"],
+                fn_kwargs={
+                    "processing_class": processing_class,
+                    "max_prompt_length": args.max_prompt_length,
+                    "max_completion_length": args.max_completion_length,
+                    # for enc-dec, we add the special tokens ([bos_token] + prompt + [eos_token]; completion + [eos_token])
+                    "add_special_tokens": False,
+                },
+                **map_kwargs,
+            )
+
+        return dataset
+
     @staticmethod
     def tokenize_row(
         features,
@@ -127,48 +193,68 @@ class AxolotlDPOTrainer(
         torch.cuda.empty_cache()
         return loss
 
-    def _get_train_sampler(self) -> Sampler | None:
+    # TODO: remove this once https://github.com/huggingface/trl/pull/3377 is in a release
+    def evaluation_loop(
+        self,
+        dataloader: DataLoader,
+        description: str,
+        prediction_loss_only: Optional[bool] = None,
+        ignore_keys: Optional[list[str]] = None,
+        metric_key_prefix: str = "eval",
+    ) -> EvalLoopOutput:
         """
-        Helper method to get the sampler for training. Handles cases for sequence
+        Overriding built-in evaluation loop to store metrics for each batch.
-        parallelism, sample packing, and curriculum sampling (sequential).
+        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.
 
-        Returns:
+        Works both with or without labels.
-            If the dataset is non-empty, a sampler is returned, the type of which
-                depends on the passed training args.
         """
-        import torch.distributed as dist
 
-        if dist.get_rank() == 0:
+        # Sample and save to game log if requested (for one batch to save time)
-            import ipdb
+        if self.generate_during_eval:
+            # Generate random indices within the range of the total number of samples
+            num_samples = len(dataloader.dataset)
+            random_indices = random.sample(
+                range(num_samples), k=self.args.eval_batch_size
+            )
 
-            ipdb.set_trace()
+            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader
-        dist.barrier()
+            random_batch_dataset = dataloader.dataset.select(random_indices)
-        if dist.get_rank() == 1:
+            random_batch = self.data_collator(random_batch_dataset)
-            import ipdb
+            random_batch = self._prepare_inputs(random_batch)
 
-            ipdb.set_trace()
+            policy_output_decoded, ref_output_decoded = (
-        dist.barrier()
+                self.generate_from_model_and_ref(self.model, random_batch)
+            )
 
-        if self.args.sequence_parallel_degree > 1:
+            table = pd.DataFrame(
-            return self._sp_get_train_sampler(self.train_dataset)
+                columns=["Prompt", "Policy", "Ref Model"],
+                data=[
+                    [prompt, pol[len(prompt) :], ref[len(prompt) :]]
+                    for prompt, pol, ref in zip(
+                        random_batch_dataset["prompt"],
+                        policy_output_decoded,
+                        ref_output_decoded,
+                    )
+                ],
+            )
+            if "wandb" in self.args.report_to and self.accelerator.is_main_process:
+                wandb.log({"game_log": wandb.Table(data=table)})
 
-        return super()._get_train_sampler()
+            if "comet_ml" in self.args.report_to:
+                log_table_to_comet_experiment(
+                    name="game_log.csv",
+                    table=table,
+                )
 
-    def _get_eval_sampler(self, eval_dataset: Dataset | None = None) -> Sampler | None:
+        # Base evaluation
-        """
+        initial_output = super(  # pylint: disable=bad-super-call
-        Helper method to get the sampler for evaluation. Handles sequence parallelism
+            DPOTrainer, self
-        and sample packing cases.
+        ).evaluation_loop(
+            dataloader,
+            description,
+            prediction_loss_only,
+            ignore_keys,
+            metric_key_prefix,
+        )
 
-        Args:
+        return initial_output
-            eval_dataset: Evaluation dataset.
-
-        Returns:
-            If the dataset is non-empty, a sampler is returned, the type of which
-                depends on the passed training args.
-        """
-        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
-
-        if self.args.sequence_parallel_degree > 1:
-            return self._sp_get_eval_sampler(eval_dataset)
-
-        return super()._get_eval_sampler(eval_dataset)

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

@@ -266,6 +266,9 @@ class AxolotlGRPOSequenceParallelTrainer(AxolotlGRPOTrainer):
             self.accelerator.even_batches = False
 
         # Return unprepared dataloader if using sequence parallelism
+        # TODO(djsaunde): We might be able to use `accelerate`'s dataloader preparation
+        # if we use `dispatch_batches` and `slice_fn_for_dispatch` properly (i.e.,
+        # slice each batch along the sequence dimension).
         if self.args.sequence_parallel_degree > 1:
             return dataloader
 

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

@@ -20,15 +20,25 @@ 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,15 +17,25 @@ 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,11 +20,15 @@ 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
 
@@ -34,6 +38,10 @@ _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,
@@ -162,6 +170,10 @@ 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,9 +19,15 @@ 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
 
@@ -30,6 +36,10 @@ _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,12 +16,22 @@ 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,
@@ -150,6 +160,9 @@ 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,11 +19,15 @@ 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
 
@@ -31,6 +35,10 @@ _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,10 +13,16 @@ 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
 
@@ -25,6 +31,10 @@ _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,12 +14,22 @@ 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,13 +12,20 @@ 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
 
@@ -27,6 +34,10 @@ _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,6 +14,10 @@ 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,11 +13,21 @@ 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,16 +7,24 @@ 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/train.py

@@ -289,16 +289,18 @@ def save_trained_model(
                 os.remove(os.path.join(cfg.output_dir, "model.safetensors"))
             except FileNotFoundError:
                 pass
-    elif cfg.local_rank == 0:
+    else:
-        if cfg.flash_optimum and BetterTransformer:
+        if cfg.local_rank == 0:
-            model = BetterTransformer.reverse(model)
+            if cfg.flash_optimum and BetterTransformer:
+                model = BetterTransformer.reverse(model)
 
-        if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
+            if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
-            trainer.model.save_pretrained(
+                trainer.model.save_pretrained(
-                cfg.output_dir, safe_serialization=safe_serialization
+                    cfg.output_dir, safe_serialization=safe_serialization
-            )
+                )
 
-        model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
+            model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
+        trainer.accelerator.wait_for_everyone()
 
     if hasattr(cfg, "llmcompressor") and cfg.llmcompressor:
         # TODO: add integration support so this can be implemented completely within the plugin

src/axolotl/utils/ctx_managers/sequence_parallel.py

@@ -1,7 +1,6 @@
 """Module for Axolotl trainer sequence parallelism manager and utilities"""
 
 import functools
-import inspect
 
 import torch
 import torch.distributed as dist
@@ -33,7 +32,7 @@ def apply_sequence_parallelism(
     to only keep the last N tokens in the sequence during generation.
 
     Args:
-        batch: Dictionary of model arguments (e.g., input_ids, attention_mask, etc.).
+        batch: Batch dictionary (e.g., input_ids, attention_mask, etc.).
         local_rank: Local rank in the sequence parallel group.
         local_world_size: World size of the sequence parallel group.
         gradient_accumulation_steps: Number of steps to accumulate gradients over.
@@ -207,26 +206,12 @@ class SequenceParallelContextManager:
     def __enter__(self):
         # Forward pre-hook to apply sequence parallelism
         def sequence_parallel_pre_hook(_, args, kwargs):
-            # Convert all args to kwargs using the model's forward function signature
+            # Apply sequence parallelism to kwargs and get original sequence length and padding info
-            updated_kwargs = kwargs.copy()
+            kwargs, self.original_seq_len, self.pad_len = (
-
+                self.apply_sequence_parallelism(batch=kwargs)
-            # Get parameter names from the model's forward function
-            forward_params = list(
-                inspect.signature(self.models[0].forward).parameters.keys()
             )
 
-            # Map args to their parameter names
+            return args, kwargs
-            for i, arg in enumerate(args):
-                if i < len(forward_params):
-                    param_name = forward_params[i]
-                    updated_kwargs[param_name] = arg
-
-            # Apply sequence parallelism to empty args and updated kwargs
-            updated_kwargs, self.original_seq_len, self.pad_len = (
-                self.apply_sequence_parallelism(updated_kwargs)
-            )
-
-            return (), updated_kwargs
 
         # Forward post-hook to gather outputs
         def sequence_parallel_post_hook(_, __, output: ModelOutput) -> ModelOutput:

src/axolotl/utils/data/rl.py

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

src/axolotl/utils/data/sft.py

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