moving out all diff attn code to plugin repo

* basic evaluate CLI command / codepath

* tests for evaluate CLI command

* fixes and cleanup

* review comments; slightly DRYing up things

---------

Co-authored-by: Dan Saunders <danjsaund@gmail.com>

.gitignore

@@ -186,3 +186,6 @@ out/
 
 # vim
 *.swp
+
+# symlinked to axolotl-artifacts in docker containers
+outputs

cicd/cicd.sh

@@ -4,7 +4,6 @@ set -e
 python -c "import torch; assert '$PYTORCH_VERSION' in torch.__version__"
 
 pytest -v --durations=10 -n8 --ignore=tests/e2e/ --ignore=tests/patched/ /workspace/axolotl/tests/
-# pytest -v --durations=10 -n8 --dist loadfile /workspace/axolotl/tests/patched/
 pytest -v --durations=10 /workspace/axolotl/tests/e2e/patched/
 pytest -v --durations=10 /workspace/axolotl/tests/e2e/integrations/
 pytest -v --durations=10 --ignore=tests/e2e/patched/ --ignore=tests/e2e/multigpu/ --ignore=tests/e2e/integrations/ /workspace/axolotl/tests/e2e/

cicd/multigpu.py

@@ -1,6 +1,6 @@
 """
- modal application to run axolotl gpu tests in Modal
- """
+modal application to run axolotl gpu tests in Modal
+"""
 # pylint: disable=duplicate-code
 
 import os

docs/dataset-formats/pretraining.qmd

@@ -19,14 +19,7 @@ For pretraining, there is no prompt template or roles.  The only required field
 Axolotl usually loads the entire dataset into memory. This will be challenging for large datasets. Use the following config to enable streaming:
 
 ```{.yaml filename="config.yaml"}
-pretraining_dataset:
-  - name:
-    path:
-    split:
-    text_column: # column in dataset with the data, usually `text`
-    type: pretrain
-    trust_remote_code:
-    skip: # number of rows of data to skip over from the beginning
+pretraining_dataset: # hf path only
 ...
 ```
 

src/axolotl/cli/__init__.py

@@ -202,7 +202,7 @@ def do_inference(
         )
     elif cfg.chat_template:
         chat_template_str = get_chat_template(cfg.chat_template)
-    elif cfg.datasets[0].type == "chat_template":
+    elif cfg.datasets and cfg.datasets[0].type == "chat_template":
         chat_template_str = get_chat_template_from_config(
             cfg=cfg, ds_cfg=cfg.datasets[0], tokenizer=tokenizer
         )

src/axolotl/cli/evaluate.py

@@ -3,7 +3,7 @@ CLI to run training on a model
 """
 import logging
 from pathlib import Path
-from typing import Union
+from typing import Dict, Union
 
 import fire
 from dotenv import load_dotenv
@@ -23,7 +23,7 @@ from axolotl.evaluate import evaluate
 LOG = logging.getLogger("axolotl.cli.evaluate")
 
 
-def do_evaluate(cfg, cli_args) -> None:
+def do_evaluate(cfg, cli_args) -> Dict[str, float]:
     # pylint: disable=duplicate-code
     print_axolotl_text_art()
     check_accelerate_default_config()
@@ -34,7 +34,7 @@ def do_evaluate(cfg, cli_args) -> None:
     else:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
-    evaluate(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
+    return evaluate(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
 
 
 def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs) -> None:

src/axolotl/cli/main.py

@@ -1,11 +1,13 @@
 """CLI definition for various axolotl commands."""
 # pylint: disable=redefined-outer-name
+
 import subprocess  # nosec B404
 from typing import Optional
 
 import click
 
 import axolotl
+from axolotl.cli.plugins import setup_plugin_commands
 from axolotl.cli.utils import (
     add_options_from_config,
     add_options_from_dataclass,
@@ -77,6 +79,9 @@ def evaluate(config: str, accelerate: bool, **kwargs):
     """Evaluate a model."""
     kwargs = {k: v for k, v in kwargs.items() if v is not None}
 
+    # Enable expandable segments for cuda allocation to improve VRAM usage
+    set_pytorch_cuda_alloc_conf()
+
     if accelerate:
         base_cmd = ["accelerate", "launch", "-m", "axolotl.cli.evaluate"]
         if config:
@@ -254,6 +259,9 @@ def fetch(directory: str, dest: Optional[str]):
     fetch_from_github(f"{directory}/", dest)
 
 
+setup_plugin_commands(cli)
+
+
 def main():
     cli()
 

src/axolotl/cli/plugins.py

@@ -0,0 +1,36 @@
+"""Module for adding click CLI commands from axolotl plugins."""
+
+import logging
+
+import click
+
+from axolotl.cli.utils import add_options_from_config, add_options_from_dataclass
+from axolotl.logging_config import configure_logging
+from axolotl.utils.config.models.input.v0_4_1 import AxolotlInputConfig
+
+configure_logging()
+LOG = logging.getLogger(__name__)
+
+
+def setup_plugin_commands(cli: click.core.Group) -> None:
+    """
+    Setup CLI commands for available plugins.
+
+    Args:
+        cli: Click CLI object to add plugin CLI options to.
+    """
+    try:
+        from axolotl_diff_transformer.convert_diff_transformer import do_cli
+        from axolotl_diff_transformer.plugin.cli import ConvertDiffTransformerCliArgs
+
+        @cli.command()
+        @click.argument("config", type=click.Path(exists=True, path_type=str))
+        @add_options_from_dataclass(ConvertDiffTransformerCliArgs)
+        @add_options_from_config(AxolotlInputConfig)
+        def convert_diff_transformer(config: str, **kwargs):
+            """Convert model attention layers to differential attention layers."""
+            kwargs = {k: v for k, v in kwargs.items() if v is not None}
+            do_cli(config=config, **kwargs)
+
+    except ImportError as exc:
+        LOG.debug("axolotl-diff-transformer not found: %s", exc)

src/axolotl/cli/utils.py

@@ -22,11 +22,11 @@ def add_options_from_dataclass(config_class: Type[Any]):
         # Process dataclass fields in reverse order for correct option ordering
         for field in reversed(dataclasses.fields(config_class)):
             field_type = field.type
-
             if get_origin(field_type) is Union and type(None) in get_args(field_type):
                 field_type = next(
                     t for t in get_args(field_type) if not isinstance(t, NoneType)
                 )
+
             if field_type == bool:
                 field_name = field.name.replace("_", "-")
                 option_name = f"--{field_name}/--no-{field_name}"
@@ -43,6 +43,7 @@ def add_options_from_dataclass(config_class: Type[Any]):
                     default=field.default,
                     help=field.metadata.get("description"),
                 )(function)
+
         return function
 
     return decorator
@@ -54,7 +55,14 @@ def add_options_from_config(config_class: Type[BaseModel]):
     def decorator(function):
         # Process model fields in reverse order for correct option ordering
         for name, field in reversed(config_class.model_fields.items()):
-            if field.annotation == bool:
+            field_type = field.annotation
+            if get_origin(field_type) is Union and type(None) in get_args(field_type):
+                field_type = next(
+                    t for t in get_args(field_type) if not isinstance(t, NoneType)
+                )
+
+            # NOTE: defaults are handled by the pydantic model config classes.
+            if field_type == bool:
                 field_name = name.replace("_", "-")
                 option_name = f"--{field_name}/--no-{field_name}"
                 function = click.option(
@@ -65,6 +73,7 @@ def add_options_from_config(config_class: Type[BaseModel]):
                 function = click.option(
                     option_name, default=None, help=field.description
                 )(function)
+
         return function
 
     return decorator
@@ -83,6 +92,8 @@ def build_command(base_cmd: List[str], options: Dict[str, Any]) -> List[str]:
         if isinstance(value, bool):
             if value:
                 cmd.append(f"--{key}")
+            else:
+                cmd.append(f"--no{key}")
         else:
             cmd.extend([f"--{key}", str(value)])
 

src/axolotl/common/cli.py

@@ -4,22 +4,26 @@ shared module for cli specific things
 
 import logging
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import TYPE_CHECKING, Optional, Union
 
 import axolotl.monkeypatch.data.batch_dataset_fetcher  # pylint: disable=unused-import  # noqa: F401
 from axolotl.logging_config import configure_logging
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.models import load_model, load_tokenizer
 
+if TYPE_CHECKING:
+    try:
+        from axolotl_diff_transformer.plugin.cli import ConvertDiffTransformerCliArgs
+    except:  # noqa: E722 # pylint: disable=bare-except # nosec B110
+        pass
+
 configure_logging()
-LOG = logging.getLogger("axolotl.common.cli")
+LOG = logging.getLogger(__name__)
 
 
 @dataclass
 class PreprocessCliArgs:
-    """
-    dataclass representing arguments for preprocessing only
-    """
+    """dataclass with arguments for preprocessing only"""
 
     debug: bool = field(default=False)
     debug_text_only: bool = field(default=False)
@@ -30,9 +34,7 @@ class PreprocessCliArgs:
 
 @dataclass
 class TrainerCliArgs:
-    """
-    dataclass representing the various non-training arguments
-    """
+    """dataclass with various non-training arguments"""
 
     debug: bool = field(default=False)
     debug_text_only: bool = field(default=False)
@@ -45,9 +47,7 @@ class TrainerCliArgs:
 
 @dataclass
 class EvaluateCliArgs:
-    """
-    dataclass representing the various evaluation arguments
-    """
+    """dataclass with various evaluation arguments"""
 
     debug: bool = field(default=False)
     debug_text_only: bool = field(default=False)
@@ -57,7 +57,7 @@ class EvaluateCliArgs:
 def load_model_and_tokenizer(
     *,
     cfg: DictDefault,
-    cli_args: TrainerCliArgs,
+    cli_args: Union[TrainerCliArgs, EvaluateCliArgs, "ConvertDiffTransformerCliArgs"],
 ):
     LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
     tokenizer = load_tokenizer(cfg)

src/axolotl/core/trainer_builder.py

@@ -293,7 +293,7 @@ class AxolotlTrainingArguments(AxolotlTrainingMixins, TrainingArguments):
     """
     Training arguments for Causal trainer
 
-    This code is duplicated due to HF TrainingArguments not setting output_dir with a defaujlt value
+    This code is duplicated due to HF TrainingArguments not setting output_dir with a default value
     so it can't be used as a mixin.
     """
 

src/axolotl/evaluate.py

@@ -9,12 +9,11 @@ from typing import Dict, Optional
 import torch
 from accelerate.logging import get_logger
 
-from axolotl.common.cli import TrainerCliArgs
+from axolotl.common.cli import EvaluateCliArgs, load_model_and_tokenizer
 from axolotl.logging_config import configure_logging
 from axolotl.train import TrainDatasetMeta
-from axolotl.utils import set_pytorch_cuda_alloc_conf
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_processor, load_tokenizer
+from axolotl.utils.models import load_processor
 from axolotl.utils.trainer import setup_trainer
 
 project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
@@ -62,8 +61,9 @@ def evaluate_dataset(
     return metrics
 
 
+# pylint: disable=duplicate-code
 def evaluate(
-    *, cfg: DictDefault, cli_args: TrainerCliArgs, dataset_meta: TrainDatasetMeta
+    *, cfg: DictDefault, cli_args: EvaluateCliArgs, dataset_meta: TrainDatasetMeta
 ) -> Dict[str, float]:
     """
     Evaluate a model on training and validation datasets
@@ -79,16 +79,11 @@ def evaluate(
         - The tokenizer
         - Dictionary of evaluation metrics
     """
-    # pylint: disable=duplicate-code
-    # Enable expandable segments for cuda allocation to improve VRAM usage
-    set_pytorch_cuda_alloc_conf()
+    # Load model
+    LOG.debug("loading model for evaluation...")
 
-    # Load tokenizer
-    LOG.debug(
-        f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}",
-        main_process_only=True,
-    )
-    tokenizer = load_tokenizer(cfg)
+    model, tokenizer = load_model_and_tokenizer(cfg=cfg, cli_args=cli_args)
+    model = model.to(cfg.device, dtype=cfg.torch_dtype)
 
     # Load processor for multimodal models if needed
     processor = None
@@ -100,12 +95,6 @@ def evaluate(
     eval_dataset = dataset_meta.eval_dataset
     total_num_steps = dataset_meta.total_num_steps
 
-    # Load model
-    LOG.debug("loading model for evaluation...")
-    model, _ = load_model(
-        cfg, tokenizer, processor=processor, inference=cli_args.inference
-    )
-
     # Set up trainer
     trainer = setup_trainer(
         cfg,

src/axolotl/integrations/config.py

@@ -43,10 +43,12 @@ def merge_input_args():
     input_args: List[str] = plugin_manager.get_input_args()
     plugin_classes = []
     dynamic_input = ""
+
     for plugin_args in input_args:
         plugin_module, plugin_cls = plugin_args.rsplit(".", 1)
         dynamic_input += f"from {plugin_module} import {plugin_cls}\n"
         plugin_classes.append(plugin_cls)
+
     if dynamic_input:
         dynamic_input += f"class AxolotlConfigWCapabilities(AxolotlConfigWCapabilitiesBase, {', '.join(plugin_classes)}):\n    pass\n"
         dynamic_input += f"class AxolotlInputConfig(AxolotlInputConfigBase, {', '.join(plugin_classes)}):\n    pass\n"
@@ -62,4 +64,5 @@ def merge_input_args():
             "AxolotlConfigWCapabilities"
         ]
         return AxolotlConfigWCapabilities, AxolotlInputConfig
+
     return AxolotlConfigWCapabilitiesBase, AxolotlInputConfigBase

src/axolotl/utils/config/models/input/v0_4_1/__init__.py

@@ -129,7 +129,6 @@ class PretrainingDataset(BaseModel):
     type: Optional[str] = "pretrain"
     trust_remote_code: Optional[bool] = False
     data_files: Optional[str] = None
-    skip: Optional[int] = None
 
 
 class UserDefinedPrompterType(BaseModel):
@@ -368,13 +367,6 @@ class LoraConfig(BaseModel):
             loraplus_lr_embedding = float(loraplus_lr_embedding)
         return loraplus_lr_embedding
 
-    @model_validator(mode="before")
-    @classmethod
-    def validate_lora_dropout(cls, data):
-        if data.get("adapter") is not None and data.get("lora_dropout") is None:
-            data["lora_dropout"] = 0.0
-        return data
-
 
 class ReLoRAConfig(BaseModel):
     """ReLoRA configuration subset"""

src/axolotl/utils/data/sft.py

@@ -89,13 +89,11 @@ def prepare_dataset(cfg, tokenizer, processor=None):
         split = "train"
         name = None
         data_files = None
-        skip = 0
         if isinstance(cfg.pretraining_dataset, list) and isinstance(
             cfg.pretraining_dataset[0], dict
         ):
             path = cfg.pretraining_dataset[0]["path"]
             name = cfg.pretraining_dataset[0]["name"]
-            skip = cfg.pretraining_dataset[0]["skip"]
             if "split" in cfg.pretraining_dataset[0]:
                 split = cfg.pretraining_dataset[0]["split"]
 
@@ -109,14 +107,10 @@ def prepare_dataset(cfg, tokenizer, processor=None):
             cfg.pretraining_dataset[0]["type"] or "pretrain",
         )
 
-        iter_ds = load_dataset(
-            path, streaming=True, split=split, name=name, data_files=data_files
-        )
-        if skip:
-            LOG.info(f"Skipping {skip} samples from the dataset")
-            iter_ds = iter_ds.skip(skip)
         train_dataset = wrap_pretraining_dataset(
-            iter_ds,
+            load_dataset(
+                path, streaming=True, split=split, name=name, data_files=data_files
+            ),
             tokenizer,
             cfg,
             ds_wrapper_partial,

src/axolotl/utils/models.py

@@ -713,19 +713,45 @@ class ModelLoader:
         if self.cfg.flash_attention:
             if not self.cfg.sample_packing and self.cfg.s2_attention:
                 pass
-            self.model_kwargs["attn_implementation"] = "flash_attention_2"
-            self.model_config._attn_implementation = (  # pylint: disable=protected-access
-                "flash_attention_2"
-            )
+
+            if self.cfg.diff_attention:
+                self.model_kwargs[
+                    "attn_implementation"
+                ] = "differential_flash_attention_2"
+                self.model_config._attn_implementation = (  # pylint: disable=protected-access
+                    "differential_flash_attention_2"
+                )
+            else:
+                self.model_kwargs["attn_implementation"] = "flash_attention_2"
+                self.model_config._attn_implementation = (  # pylint: disable=protected-access
+                    "flash_attention_2"
+                )
         elif self.cfg.sdp_attention:
-            self.model_kwargs["attn_implementation"] = "sdpa"
-            self.model_config._attn_implementation = (  # pylint: disable=protected-access
-                "sdpa"
-            )
+            if self.cfg.diff_attention:
+                self.model_kwargs["attn_implementation"] = "differential_sdpa"
+                self.model_config._attn_implementation = (  # pylint: disable=protected-access
+                    "differential_sdpa"
+                )
+            else:
+                self.model_kwargs["attn_implementation"] = "sdpa"
+                self.model_config._attn_implementation = (  # pylint: disable=protected-access
+                    "sdpa"
+                )
         elif self.cfg.eager_attention:
-            self.model_kwargs["attn_implementation"] = "eager"
+            if self.cfg.diff_attention:
+                self.model_kwargs["attn_implementation"] = "differential_eager"
+                self.model_config._attn_implementation = (  # pylint: disable=protected-access
+                    "differential_eager"
+                )
+            else:
+                self.model_kwargs["attn_implementation"] = "eager"
+                self.model_config._attn_implementation = (  # pylint: disable=protected-access
+                    "eager"
+                )
+        elif self.cfg.diff_attention:
+            self.model_kwargs["attn_implementation"] = "differential_eager"
             self.model_config._attn_implementation = (  # pylint: disable=protected-access
-                "eager"
+                "differential_eager"
             )
 
         if self.cfg.low_cpu_mem_usage:
@@ -816,6 +842,7 @@ class ModelLoader:
 
             if self.cfg.is_multimodal:
                 self.model_config.text_config = self.text_model_config
+
             self.model = self.AutoModelLoader.from_pretrained(
                 self.base_model,
                 config=self.model_config,

src/axolotl/utils/yaml.py

@@ -0,0 +1,157 @@
+"""Utilities for YAML files."""
+
+from collections import OrderedDict
+from typing import Any, Dict, List, Set, Tuple, Union
+
+import yaml
+
+
+class YAMLOrderTracker:
+    """Tracks the order of keys and section breaks in YAML files."""
+
+    def __init__(self, yaml_path: str):
+        self.yaml_path = yaml_path
+        self.structure, self.needs_break = self._parse_yaml_structure()
+
+    def _get_indentation_level(self, line: str) -> int:
+        """Get the indentation level of a line."""
+        return len(line) - len(line.lstrip())
+
+    def _parse_yaml_structure(
+        self,
+    ) -> Tuple[Dict[str, Union[List[str], Dict]], Set[str]]:
+        """Parse the YAML file to extract structure and identify section breaks."""
+        with open(self.yaml_path, "r", encoding="utf-8") as file:
+            contents = file.readlines()
+
+        structure: OrderedDict = OrderedDict()
+        needs_break = set()  # Track which keys should have a break before them
+        current_path = []
+        last_indentation = -1
+        had_empty_line = False
+
+        for line in contents:
+            # Track empty lines and comments
+            if not line.strip() or line.strip().startswith("#"):
+                had_empty_line = True
+                continue
+
+            # Get indentation level and content
+            indentation = self._get_indentation_level(line)
+            content = line.strip()
+
+            # Skip lines that don't define keys
+            if ":" not in content:
+                continue
+
+            # Extract key
+            key = content.split(":")[0].strip()
+
+            # If this is a top-level key and we had an empty line, mark it
+            if indentation == 0:
+                if had_empty_line:
+                    needs_break.add(key)
+                had_empty_line = False
+
+            # Handle indentation changes
+            if indentation > last_indentation:
+                current_path.append(key)
+            elif indentation < last_indentation:
+                levels_up = (last_indentation - indentation) // 2
+                current_path = current_path[:-levels_up]
+                current_path[-1] = key
+            else:
+                if current_path:
+                    current_path[-1] = key
+
+            # Update structure
+            current_dict = structure
+            for path_key in current_path[:-1]:
+                if path_key not in current_dict:
+                    current_dict[path_key] = OrderedDict()
+                current_dict = current_dict[path_key]
+
+            if current_path:
+                if current_path[-1] not in current_dict:
+                    current_dict[current_path[-1]] = OrderedDict()
+
+            last_indentation = indentation
+
+        return structure, needs_break
+
+
+class OrderedDumper(yaml.SafeDumper):
+    """Custom YAML dumper that maintains dictionary order."""
+
+
+def represent_none(self, _):
+    """Represent None values as empty fields."""
+    return self.represent_scalar("tag:yaml.org,2002:null", "")
+
+
+def ordered_dict_representer(dumper: OrderedDumper, data: Dict) -> Any:
+    """Custom representer for dictionaries that maintains order."""
+    return dumper.represent_mapping("tag:yaml.org,2002:map", data.items())
+
+
+def reorder_dict(data: Dict, reference_structure: Dict) -> OrderedDict:
+    """Reorder a dictionary based on a reference structure."""
+    ordered = OrderedDict()
+
+    # First add keys that are in the reference order
+    for key in reference_structure:
+        if key in data:
+            if isinstance(reference_structure[key], dict) and isinstance(
+                data[key], dict
+            ):
+                ordered[key] = reorder_dict(data[key], reference_structure[key])
+            else:
+                ordered[key] = data[key]
+
+    # Then add any remaining keys that weren't in the reference
+    for key in data:
+        if key not in ordered:
+            ordered[key] = data[key]
+
+    return ordered
+
+
+def dump_yaml_preserved_order(
+    data: Dict, reference_yaml_path: str, output_path: str
+) -> None:
+    """Dump YAML file while preserving nested order and normalized spacing."""
+    # Get reference structure and spacing
+    tracker = YAMLOrderTracker(reference_yaml_path)
+
+    # Reorder the data
+    ordered_data = reorder_dict(data, tracker.structure)
+
+    # Register the custom representers
+    OrderedDumper.add_representer(type(None), represent_none)
+    OrderedDumper.add_representer(dict, ordered_dict_representer)
+    OrderedDumper.add_representer(OrderedDict, ordered_dict_representer)
+
+    # First dump to string
+    yaml_str = yaml.dump(
+        ordered_data, Dumper=OrderedDumper, sort_keys=False, default_flow_style=False
+    )
+
+    # Add spacing according to reference
+    lines = yaml_str.split("\n")
+    result_lines: List[str] = []
+    current_line = 0
+
+    while current_line < len(lines):
+        line = lines[current_line]
+        if line.strip() and ":" in line and not line.startswith(" "):  # Top-level key
+            key = line.split(":")[0].strip()
+            if key in tracker.needs_break:
+                # Add single empty line before this key
+                if result_lines and result_lines[-1] != "":
+                    result_lines.append("")
+        result_lines.append(line)
+        current_line += 1
+
+    # Write the final result
+    with open(output_path, "w", encoding="utf-8") as file:
+        file.write("\n".join(result_lines))

tests/cli/conftest.py

@@ -1,4 +1,5 @@
 """Shared pytest fixtures for cli module."""
+
 import pytest
 from click.testing import CliRunner
 

tests/cli/test_cli_base.py

@@ -43,14 +43,12 @@ class BaseCliTest:
             result = cli_runner.invoke(cli, [command, str(config_path)])
 
             assert mock.called
-            assert mock.call_args.args[0] == [
+            assert mock.call_args.args[0][:5] == [
                 "accelerate",
                 "launch",
                 "-m",
                 f"axolotl.cli.{command}",
                 str(config_path),
-                "--debug-num-examples",
-                "0",
             ]
             assert mock.call_args.kwargs == {"check": True}
             assert result.exit_code == 0

tests/cli/test_cli_fetch.py

@@ -1,4 +1,5 @@
 """pytest tests for axolotl CLI fetch command."""
+
 from unittest.mock import patch
 
 from axolotl.cli.main import fetch

tests/cli/test_cli_inference.py

@@ -1,4 +1,5 @@
 """pytest tests for axolotl CLI inference command."""
+
 from unittest.mock import patch
 
 from axolotl.cli.main import cli

tests/cli/test_cli_interface.py

@@ -1,4 +1,5 @@
 """General pytest tests for axolotl.cli.main interface."""
+
 from axolotl.cli.main import build_command, cli
 
 
@@ -22,6 +23,7 @@ def test_build_command():
         "--batch-size",
         "8",
         "--debug",
+        "--nouse-fp16",
     ]
 
 

tests/cli/test_cli_merge_lora.py

@@ -1,4 +1,5 @@
 """pytest tests for axolotl CLI merge_lora command."""
+
 from unittest.mock import patch
 
 from axolotl.cli.main import cli

tests/cli/test_cli_merge_sharded_fsdp_weights.py

@@ -1,5 +1,6 @@
 """pytest tests for axolotl CLI merge_sharded_fsdp_weights command."""
 # pylint: disable=duplicate-code
+
 from unittest.mock import patch
 
 from axolotl.cli.main import cli

tests/cli/test_cli_preprocess.py

@@ -1,4 +1,5 @@
 """pytest tests for axolotl CLI preprocess command."""
+
 import shutil
 from pathlib import Path
 from unittest.mock import patch

tests/cli/test_cli_shard.py

@@ -1,5 +1,6 @@
 """pytest tests for axolotl CLI shard command."""
 # pylint: disable=duplicate-code
+
 from unittest.mock import patch
 
 from axolotl.cli.main import cli
@@ -11,14 +12,12 @@ def test_shard_with_accelerate(cli_runner, config_path):
         result = cli_runner.invoke(cli, ["shard", str(config_path), "--accelerate"])
 
         assert mock.called
-        assert mock.call_args.args[0] == [
+        assert mock.call_args.args[0][:5] == [
             "accelerate",
             "launch",
             "-m",
             "axolotl.cli.shard",
             str(config_path),
-            "--debug-num-examples",
-            "0",
         ]
         assert mock.call_args.kwargs == {"check": True}
         assert result.exit_code == 0

tests/cli/test_cli_version.py

@@ -1,4 +1,5 @@
 """pytest tests for axolotl CLI --version"""
+
 from axolotl.cli.main import cli
 
 

tests/cli/test_utils.py

@@ -1,5 +1,6 @@
 """pytest tests for axolotl CLI utils."""
 # pylint: disable=redefined-outer-name
+
 import json
 from unittest.mock import Mock, patch
 

tests/e2e/integrations/test_cut_cross_entropy.py

@@ -2,6 +2,8 @@
 Simple end-to-end test for Cut Cross Entropy integration
 """
 
+from pathlib import Path
+
 import pytest
 
 from axolotl.cli import load_datasets
@@ -11,8 +13,6 @@ from axolotl.utils import get_pytorch_version
 from axolotl.utils.config import normalize_config, prepare_plugins
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists
-
 # pylint: disable=duplicate-code
 
 
@@ -67,7 +67,7 @@ class TestCutCrossEntropyIntegration:
                 train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
         else:
             train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-            check_model_output_exists(temp_dir, cfg)
+            assert (Path(temp_dir) / "model.safetensors").exists()
 
     @pytest.mark.parametrize(
         "attention_type",
@@ -95,4 +95,4 @@ class TestCutCrossEntropyIntegration:
                 train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
         else:
             train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-            check_model_output_exists(temp_dir, cfg)
+            assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/integrations/test_liger.py

@@ -1,6 +1,7 @@
 """
 Simple end-to-end test for Liger integration
 """
+from pathlib import Path
 
 from e2e.utils import require_torch_2_4_1
 
@@ -10,8 +11,6 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config, prepare_plugins
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists
-
 
 class LigerIntegrationTestCase:
     """
@@ -61,7 +60,7 @@ class LigerIntegrationTestCase:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()
 
     @require_torch_2_4_1
     def test_llama_w_flce(self, temp_dir):
@@ -106,4 +105,4 @@ class LigerIntegrationTestCase:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/patched/test_4d_multipack_llama.py

@@ -5,6 +5,7 @@ E2E tests for multipack fft llama using 4d attention masks
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, require_torch_2_3_1, with_temp_dir
+from ..utils import require_torch_2_3_1, with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -66,7 +67,7 @@ class Test4dMultipackLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_torch_lora_packing(self, temp_dir):
@@ -110,4 +111,4 @@ class Test4dMultipackLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()

tests/e2e/patched/test_fa_xentropy.py

@@ -4,6 +4,7 @@ E2E tests for lora llama
 
 import logging
 import os
+from pathlib import Path
 
 import pytest
 from transformers.utils import is_torch_bf16_gpu_available
@@ -14,7 +15,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, check_tensorboard
+from ..utils import check_tensorboard
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -81,7 +82,7 @@ class TestFAXentropyLlama:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 1.5, "Train Loss is too high"

tests/e2e/patched/test_falcon_samplepack.py

@@ -5,6 +5,7 @@ E2E tests for falcon
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -68,7 +69,7 @@ class TestFalconPatched(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_ft(self, temp_dir):
@@ -108,4 +109,4 @@ class TestFalconPatched(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/patched/test_fused_llama.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 import pytest
 from transformers.utils import is_torch_bf16_gpu_available
@@ -15,7 +16,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -72,4 +73,4 @@ class TestFusedLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/patched/test_llama_s2_attention.py

@@ -5,6 +5,7 @@ E2E tests for llama w/ S2 attn
 import logging
 import os
 import unittest
+from pathlib import Path
 
 import pytest
 
@@ -14,7 +15,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -70,7 +71,7 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_fft_s2_attn(self, temp_dir):
@@ -110,4 +111,4 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/patched/test_lora_llama_multipack.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 import pytest
 from transformers.utils import is_auto_gptq_available, is_torch_bf16_gpu_available
@@ -15,7 +16,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -75,7 +76,7 @@ class TestLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @pytest.mark.skipif(not is_auto_gptq_available(), reason="auto-gptq not available")
     @with_temp_dir
@@ -125,4 +126,4 @@ class TestLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()

tests/e2e/patched/test_mistral_samplepack.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -68,7 +69,7 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_ft_packing(self, temp_dir):
@@ -109,4 +110,4 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/patched/test_mixtral_samplepack.py

@@ -5,6 +5,7 @@ E2E tests for mixtral
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -65,7 +66,7 @@ class TestMixtral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_ft(self, temp_dir):
@@ -107,4 +108,4 @@ class TestMixtral(unittest.TestCase):
             "MixtralFlashAttention2"
             in model.model.layers[0].self_attn.__class__.__name__
         )
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/patched/test_phi_multipack.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, with_temp_dir
+from ..utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -68,7 +69,7 @@ class TestPhiMultipack(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()
 
     @with_temp_dir
     def test_qlora_packed(self, temp_dir):
@@ -119,4 +120,4 @@ class TestPhiMultipack(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()

tests/e2e/patched/test_resume.py

@@ -6,6 +6,7 @@ import logging
 import os
 import re
 import subprocess
+from pathlib import Path
 
 from transformers.utils import is_torch_bf16_gpu_available
 
@@ -15,7 +16,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, most_recent_subdir
+from ..utils import most_recent_subdir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -82,7 +83,7 @@ class TestResumeLlama:
         cli_args = TrainerCliArgs()
 
         train(cfg=resume_cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
         tb_log_path_1 = most_recent_subdir(temp_dir + "/runs")
         cmd = f"tensorboard --inspect  --logdir {tb_log_path_1}"

tests/e2e/patched/test_unsloth_qlora.py

@@ -3,6 +3,7 @@ e2e tests for unsloth qlora
 """
 import logging
 import os
+from pathlib import Path
 
 import pytest
 
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from ..utils import check_model_output_exists, check_tensorboard
+from ..utils import check_tensorboard
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -76,7 +77,7 @@ class TestUnslothQLoRA:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss is too high"
@@ -126,7 +127,7 @@ class TestUnslothQLoRA:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss is too high"
@@ -181,7 +182,7 @@ class TestUnslothQLoRA:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 2.0, "Train Loss is too high"

tests/e2e/test_dpo.py

@@ -15,7 +15,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -68,7 +68,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
 
     @with_temp_dir
     def test_dpo_nll_lora(self, temp_dir):
@@ -113,7 +113,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
 
     @with_temp_dir
     def test_dpo_use_weighting(self, temp_dir):
@@ -158,7 +158,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
 
     @pytest.mark.skip("kto_pair no longer supported in trl")
     @with_temp_dir
@@ -203,7 +203,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
 
     @with_temp_dir
     def test_ipo_lora(self, temp_dir):
@@ -247,7 +247,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
 
     @with_temp_dir
     def test_orpo_lora(self, temp_dir):
@@ -294,7 +294,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()
 
     @pytest.mark.skip(reason="Fix the implementation")
     @with_temp_dir
@@ -358,4 +358,4 @@ class TestDPOLlamaLora(unittest.TestCase):
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-20", cfg)
+        assert (Path(temp_dir) / "checkpoint-20/adapter_model.safetensors").exists()

tests/e2e/test_embeddings_lr.py

@@ -5,6 +5,7 @@ E2E tests for llama pretrain
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
+from .utils import check_tensorboard, with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -61,7 +62,7 @@ class TestEmbeddingsLrScale(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 2.0, "Loss is too high"
@@ -105,7 +106,7 @@ class TestEmbeddingsLrScale(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/train_loss", 2.0, "Loss is too high"

tests/e2e/test_falcon.py

@@ -5,6 +5,7 @@ E2E tests for falcon
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -70,7 +71,7 @@ class TestFalcon(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_lora_added_vocab(self, temp_dir):
@@ -123,7 +124,7 @@ class TestFalcon(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_ft(self, temp_dir):
@@ -162,4 +163,4 @@ class TestFalcon(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/test_llama.py

@@ -4,8 +4,7 @@ E2E tests for llama
 
 import logging
 import os
-
-from e2e.utils import check_model_output_exists
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -61,7 +60,7 @@ class TestLlama:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()
 
     def test_fix_untrained_tokens(self, temp_dir):
         # pylint: disable=duplicate-code
@@ -104,7 +103,7 @@ class TestLlama:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()
 
     def test_batch_flattening(self, temp_dir):
         # pylint: disable=duplicate-code
@@ -143,4 +142,4 @@ class TestLlama:
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/test_llama_pretrain.py

@@ -5,6 +5,7 @@ E2E tests for llama pretrain
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -63,4 +64,4 @@ class TestPretrainLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/test_llama_vision.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -67,7 +68,7 @@ class TestLlamaVision(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
 
     @with_temp_dir
     def test_lora_llama_vision_multimodal_dataset(self, temp_dir):
@@ -112,4 +113,4 @@ class TestLlamaVision(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()

tests/e2e/test_lora_llama.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -64,4 +65,4 @@ class TestLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()

tests/e2e/test_mamba.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 import pytest
 
@@ -14,7 +15,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -64,4 +65,4 @@ class TestMamba(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/test_mistral.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from transformers.utils import is_torch_bf16_gpu_available
 
@@ -14,7 +15,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -68,7 +69,7 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_ft(self, temp_dir):
@@ -111,4 +112,4 @@ class TestMistral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/test_mixtral.py

@@ -5,6 +5,7 @@ E2E tests for mixtral
 import logging
 import os
 import unittest
+from pathlib import Path
 
 import torch
 from transformers.utils import is_torch_bf16_gpu_available
@@ -15,7 +16,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -78,7 +79,7 @@ class TestMixtral(unittest.TestCase):
             model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
             == torch.float32
         )
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_qlora_wo_fa2(self, temp_dir):
@@ -132,7 +133,7 @@ class TestMixtral(unittest.TestCase):
             model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
             == torch.float32
         )
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_16bit_lora_w_fa2(self, temp_dir):
@@ -189,7 +190,7 @@ class TestMixtral(unittest.TestCase):
             model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
             == torch.float32
         )
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_16bit_lora_wo_fa2(self, temp_dir):
@@ -246,7 +247,7 @@ class TestMixtral(unittest.TestCase):
             model.base_model.model.model.layers[0].block_sparse_moe.gate.weight.dtype
             == torch.float32
         )
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_ft(self, temp_dir):
@@ -286,4 +287,4 @@ class TestMixtral(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()

tests/e2e/test_optimizers.py

@@ -5,6 +5,7 @@ E2E tests for custom optimizers using Llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, require_torch_2_5_1, with_temp_dir
+from .utils import require_torch_2_5_1, with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -64,7 +65,7 @@ class TestCustomOptimizers(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     @require_torch_2_5_1
@@ -108,7 +109,7 @@ class TestCustomOptimizers(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()
 
     @with_temp_dir
     def test_fft_schedule_free_adamw(self, temp_dir):
@@ -144,4 +145,4 @@ class TestCustomOptimizers(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/test_phi.py

@@ -5,6 +5,7 @@ E2E tests for lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -66,7 +67,7 @@ class TestPhi(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "pytorch_model.bin").exists()
 
     @with_temp_dir
     def test_phi_qlora(self, temp_dir):
@@ -115,4 +116,4 @@ class TestPhi(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()

tests/e2e/test_relora_llama.py

@@ -13,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
+from .utils import check_tensorboard, with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -78,10 +78,10 @@ class TestReLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(Path(temp_dir) / "checkpoint-100/adapter", cfg)
         assert (
-            Path(temp_dir) / "checkpoint-100/relora/model.safetensors"
-        ).exists(), "Relora model checkpoint not found"
+            Path(temp_dir) / "checkpoint-100/adapter/adapter_model.safetensors"
+        ).exists()
+        assert (Path(temp_dir) / "checkpoint-100/relora/model.safetensors").exists()
 
         check_tensorboard(
             temp_dir + "/runs", "train/grad_norm", 0.2, "grad_norm is too high"

tests/e2e/test_reward_model_llama.py

@@ -5,6 +5,7 @@ E2E tests for reward model lora llama
 import logging
 import os
 import unittest
+from pathlib import Path
 
 from axolotl.cli import load_datasets
 from axolotl.common.cli import TrainerCliArgs
@@ -12,7 +13,7 @@ from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
 
-from .utils import check_model_output_exists, with_temp_dir
+from .utils import with_temp_dir
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
@@ -70,4 +71,4 @@ class TestRewardModelLoraLlama(unittest.TestCase):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
         train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
+        assert (Path(temp_dir) / "adapter_model.bin").exists()

tests/e2e/utils.py

@@ -14,8 +14,6 @@ import torch
 from packaging import version
 from tbparse import SummaryReader
 
-from axolotl.utils.dict import DictDefault
-
 
 def with_temp_dir(test_func):
     @wraps(test_func)
@@ -95,27 +93,3 @@ def check_tensorboard(
     df = reader.scalars  # pylint: disable=invalid-name
     df = df[(df.tag == tag)]  # pylint: disable=invalid-name
     assert df.value.values[-1] < lt_val, assertion_err
-
-
-def check_model_output_exists(temp_dir: str, cfg: DictDefault) -> None:
-    """
-    helper function to check if a model output file exists after training
-
-    checks based on adapter or not and if safetensors saves are enabled or not
-    """
-
-    if cfg.save_safetensors:
-        if not cfg.adapter:
-            assert (Path(temp_dir) / "model.safetensors").exists()
-        else:
-            assert (Path(temp_dir) / "adapter_model.safetensors").exists()
-    else:
-        # check for both, b/c in trl, it often defaults to saving safetensors
-        if not cfg.adapter:
-            assert (Path(temp_dir) / "pytorch_model.bin").exists() or (
-                Path(temp_dir) / "model.safetensors"
-            ).exists()
-        else:
-            assert (Path(temp_dir) / "adapter_model.bin").exists() or (
-                Path(temp_dir) / "adapter_model.safetensors"
-            ).exists()

tests/test_lora.py

@@ -1,69 +0,0 @@
-"""
-tests for loading loras
-"""
-from axolotl.utils.config import normalize_config, validate_config
-from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_tokenizer
-
-# pylint: disable=duplicate-code
-minimal_config = DictDefault(
-    {
-        "base_model": "HuggingFaceTB/SmolLM2-135M",
-        "learning_rate": 0.000001,
-        "datasets": [
-            {
-                "path": "mhenrichsen/alpaca_2k_test",
-                "type": "alpaca",
-            }
-        ],
-        "micro_batch_size": 1,
-        "gradient_accumulation_steps": 1,
-    }
-)
-
-
-class TestLoRALoad:
-    """
-    Test class for loading LoRA weights
-    """
-
-    def test_load_lora_weights(self):
-        cfg = DictDefault(
-            {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "adapter": "lora",
-                "lora_r": 8,
-                "lora_alpha": 16,
-                "lora_dropout": 0.0,
-                "lora_target_linear": True,
-                "micro_batch_size": 1,
-                "gradient_accumulation_steps": 1,
-                "sequence_len": 1024,
-            }
-            | minimal_config
-        )
-        cfg = validate_config(cfg)
-        normalize_config(cfg)
-        tokenizer = load_tokenizer(cfg)
-        load_model(cfg, tokenizer)
-
-    def test_load_lora_weights_empty_dropout(self):
-        cfg = DictDefault(
-            {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "adapter": "lora",
-                "lora_r": 8,
-                "lora_alpha": 16,
-                "lora_dropout": None,
-                "lora_target_linear": True,
-                "micro_batch_size": 1,
-                "gradient_accumulation_steps": 1,
-                "sequence_len": 1024,
-            }
-            | minimal_config
-        )
-        cfg = validate_config(cfg)
-        normalize_config(cfg)
-        assert cfg.lora_dropout == 0.0
-        tokenizer = load_tokenizer(cfg)
-        load_model(cfg, tokenizer)