add support for MoRA

* set chat_template in datasets config automatically

* dynamic chat_template, not jsut chatml

.github/workflows/base.yml

@@ -30,7 +30,7 @@ jobs:
           - cuda: "121"
             cuda_version: 12.1.0
             python_version: "3.11"
-            pytorch: 2.2.1
+            pytorch: 2.2.2
             torch_cuda_arch_list: "7.0 7.5 8.0 8.6 8.7 8.9 9.0+PTX"
           - cuda: "121"
             cuda_version: 12.1.0

.github/workflows/main.yml

@@ -28,7 +28,7 @@ jobs:
           - cuda: 121
             cuda_version: 12.1.0
             python_version: "3.11"
-            pytorch: 2.2.1
+            pytorch: 2.2.2
             axolotl_extras:
           - cuda: 121
             cuda_version: 12.1.0
@@ -89,7 +89,7 @@ jobs:
           - cuda: 121
             cuda_version: 12.1.0
             python_version: "3.11"
-            pytorch: 2.2.1
+            pytorch: 2.2.2
             axolotl_extras:
           - cuda: 121
             cuda_version: 12.1.0
@@ -125,3 +125,45 @@ jobs:
              ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
              ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
           labels: ${{ steps.metadata.outputs.labels }}
+
+  build-axolotl-cloud-no-tmux:
+    needs: build-axolotl
+    if: ${{ ! contains(github.event.commits[0].message, '[skip docker]]') && github.repository_owner == 'OpenAccess-AI-Collective' }}
+    # this job needs to be run on self-hosted GPU runners...
+    strategy:
+      matrix:
+        include:
+          - cuda: 121
+            cuda_version: 12.1.0
+            python_version: "3.11"
+            pytorch: 2.3.0
+            axolotl_extras:
+    runs-on: axolotl-gpu-runner
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Docker metadata
+        id: metadata
+        uses: docker/metadata-action@v5
+        with:
+          images: winglian/axolotl-cloud-term
+      - name: Login to Docker Hub
+        uses: docker/login-action@v3
+        with:
+          username: ${{ secrets.DOCKERHUB_USERNAME }}
+          password: ${{ secrets.DOCKERHUB_TOKEN }}
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@v2
+      - name: Build
+        uses: docker/build-push-action@v5
+        with:
+          context: .
+          build-args: |
+            BASE_TAG=${{ github.ref_name }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+            CUDA=${{ matrix.cuda }}
+          file: ./docker/Dockerfile-cloud-no-tmux
+          push: ${{ github.event_name != 'pull_request' }}
+          tags: |
+             ${{ steps.metadata.outputs.tags }}-py${{ matrix.python_version }}-cu${{ matrix.cuda }}-${{ matrix.pytorch }}${{ matrix.axolotl_extras != '' && '-' || '' }}${{ matrix.axolotl_extras }}
+             ${{ (matrix.is_latest) && format('{0}-latest', steps.metadata.outputs.tags) || '' }}
+          labels: ${{ steps.metadata.outputs.labels }}

.github/workflows/nightlies.yml

@@ -27,7 +27,7 @@ jobs:
           - cuda: 121
             cuda_version: 12.1.0
             python_version: "3.11"
-            pytorch: 2.2.1
+            pytorch: 2.2.2
             axolotl_extras:
           - cuda: 121
             cuda_version: 12.1.0
@@ -89,7 +89,7 @@ jobs:
           - cuda: 121
             cuda_version: 12.1.0
             python_version: "3.11"
-            pytorch: 2.2.1
+            pytorch: 2.2.2
             axolotl_extras:
           - cuda: 121
             cuda_version: 12.1.0

.github/workflows/tests.yml

@@ -82,7 +82,12 @@ jobs:
           - cuda: 121
             cuda_version: 12.1.0
             python_version: "3.11"
-            pytorch: 2.2.1
+            pytorch: 2.2.2
+            num_gpus: 1
+          - cuda: 121
+            cuda_version: 12.1.0
+            python_version: "3.11"
+            pytorch: 2.3.0
             num_gpus: 1
     steps:
       - name: Checkout

README.md

@@ -34,6 +34,7 @@ Features:
   - [Mac](#mac)
   - [Google Colab](#google-colab)
   - [Launching on public clouds via SkyPilot](#launching-on-public-clouds-via-skypilot)
+  - [Launching on public clouds via dstack](#launching-on-public-clouds-via-dstack)
 - [Dataset](#dataset)
 - [Config](#config)
   - [Train](#train)
@@ -123,11 +124,11 @@ accelerate launch -m axolotl.cli.train examples/openllama-3b/lora.yml
 
 # inference
 accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
-    --lora_model_dir="./lora-out"
+    --lora_model_dir="./outputs/lora-out"
 
 # gradio
 accelerate launch -m axolotl.cli.inference examples/openllama-3b/lora.yml \
-    --lora_model_dir="./lora-out" --gradio
+    --lora_model_dir="./outputs/lora-out" --gradio
 
 # remote yaml files - the yaml config can be hosted on a public URL
 # Note: the yaml config must directly link to the **raw** yaml
@@ -292,6 +293,42 @@ HF_TOKEN=xx sky launch axolotl.yaml --env HF_TOKEN
 HF_TOKEN=xx BUCKET=<unique-name> sky spot launch axolotl-spot.yaml --env HF_TOKEN --env BUCKET
 ```
 
+#### Launching on public clouds via dstack
+To launch on GPU instance (both on-demand and spot instances) on public clouds (GCP, AWS, Azure, Lambda Labs, TensorDock, Vast.ai, and CUDO), you can use [dstack](https://dstack.ai/).
+
+Write a job description in YAML as below:
+
+```yaml
+# dstack.yaml
+type: task
+
+image: winglian/axolotl-cloud:main-20240429-py3.11-cu121-2.2.2
+
+env:
+  - HUGGING_FACE_HUB_TOKEN
+  - WANDB_API_KEY
+
+commands:
+  - accelerate launch -m axolotl.cli.train config.yaml
+
+ports:
+  - 6006
+
+resources:
+  gpu:
+    memory: 24GB..
+    count: 2
+```
+
+then, simply run the job with `dstack run` command. Append `--spot` option if you want spot instance. `dstack run` command will show you the instance with cheapest price across multi cloud services:
+
+```bash
+pip install dstack
+HUGGING_FACE_HUB_TOKEN=xxx WANDB_API_KEY=xxx dstack run . -f dstack.yaml # --spot
+```
+
+For further and fine-grained use cases, please refer to the official [dstack documents](https://dstack.ai/docs/) and the detailed description of [axolotl example](https://github.com/dstackai/dstack/tree/master/examples/fine-tuning/axolotl) on the official repository.
+
 ### Dataset
 
 Axolotl supports a variety of dataset formats.  It is recommended to use a JSONL.  The schema of the JSONL depends upon the task and the prompt template you wish to use.  Instead of a JSONL, you can also use a HuggingFace dataset with columns for each JSONL field.

cicd/cicd.sh

@@ -1,4 +1,5 @@
 #!/bin/bash
+set -e
 
 pytest --ignore=tests/e2e/ /workspace/axolotl/tests/
 pytest /workspace/axolotl/tests/e2e/patched/

docker/Dockerfile

@@ -11,7 +11,7 @@ ARG PYTORCH_VERSION="2.1.2"
 ENV PYTORCH_VERSION=$PYTORCH_VERSION
 
 RUN apt-get update && \
-    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev
+    apt-get install -y --allow-change-held-packages vim curl nano libnccl2 libnccl-dev rsync s3fs
 
 WORKDIR /workspace
 

docker/Dockerfile-cloud-no-tmux

@@ -0,0 +1,27 @@
+ARG BASE_TAG=main
+FROM winglian/axolotl:$BASE_TAG
+
+ENV HF_DATASETS_CACHE="/workspace/data/huggingface-cache/datasets"
+ENV HUGGINGFACE_HUB_CACHE="/workspace/data/huggingface-cache/hub"
+ENV TRANSFORMERS_CACHE="/workspace/data/huggingface-cache/hub"
+ENV HF_HOME="/workspace/data/huggingface-cache/hub"
+ENV HF_HUB_ENABLE_HF_TRANSFER="1"
+
+EXPOSE 8888
+EXPOSE 22
+
+COPY scripts/cloud-entrypoint-term.sh /root/cloud-entrypoint.sh
+COPY scripts/motd /etc/motd
+
+RUN pip install jupyterlab notebook ipywidgets && \
+    jupyter lab clean
+RUN apt install --yes --no-install-recommends openssh-server tmux sudo && \
+    pip3 install -U --no-cache-dir grpcio ray[default]==2.9.3 && \
+    mkdir -p ~/.ssh && \
+    chmod 700 ~/.ssh && \
+    printf "[ ! -z \"\$TERM\" -a -r /etc/motd ] && cat /etc/motd\n" >> ~/.bashrc && \
+    chmod +x /workspace/axolotl/scripts/cloud-entrypoint.sh && \
+    chmod +x /root/cloud-entrypoint.sh
+
+ENTRYPOINT ["/root/cloud-entrypoint.sh"]
+CMD ["sleep", "infinity"]

docs/config.qmd

@@ -186,6 +186,11 @@ eval_sample_packing:
 # The trainer will provide recommended values for these values.
 sample_packing_eff_est:
 total_num_tokens:
+# Increasing the following values helps with packing, but usually only slightly (<%1.)
+# The number of samples packed at a time.
+sample_packing_group_size: 100000
+# The number of samples which can be packed into one sequence. Increase if using a large sequence_len with many short samples.
+sample_packing_bin_size: 200
 
 # Passed through to transformers when loading the model when launched without accelerate
 # Use `sequential` when training w/ model parallelism to limit memory
@@ -285,7 +290,7 @@ lr_quadratic_warmup:
 logging_steps:
 eval_steps: # Leave empty to eval at each epoch, integers for every N steps. decimal for fraction of total steps
 evals_per_epoch: # number of times per epoch to run evals, mutually exclusive with eval_steps
-save_strategy: # Set to `no` to skip checkpoint saves
+save_strategy: # Set to `"no"` to skip checkpoint saves
 save_steps: # Leave empty to save at each epoch
 saves_per_epoch: # number of times per epoch to save a checkpoint, mutually exclusive with save_steps
 save_total_limit: # Checkpoints saved at a time

examples/cerebras/btlm-ft.yml

@@ -38,7 +38,7 @@ wandb_watch:
 wandb_name:
 wandb_log_model:
 
-output_dir: btlm-out
+output_dir: ./outputs/btlm-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 1

examples/cerebras/qlora.yml

@@ -25,7 +25,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 batch_size: 4
 micro_batch_size: 4
 num_epochs: 2

examples/code-llama/13b/lora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/code-llama/13b/qlora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/code-llama/34b/lora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/code-llama/34b/qlora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/code-llama/7b/lora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/code-llama/7b/qlora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/colab-notebooks/colab-axolotl-example.ipynb

@@ -1,216 +1,223 @@
 {
-  "cells": [
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "AKjdG7tbTb-n"
-      },
-      "source": [
-        "# Example notebook for running Axolotl on google colab"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "RcbNpOgWRcii"
-      },
-      "outputs": [],
-      "source": [
-        "import torch\n",
-        "# Check so there is a gpu available, a T4(free tier) is enough to run this notebook\n",
-        "assert (torch.cuda.is_available()==True)"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "h3nLav8oTRA5"
-      },
-      "source": [
-        "## Install Axolotl and dependencies"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "3c3yGAwnOIdi",
-        "outputId": "e3777b5a-40ef-424f-e181-62dfecd1dd01"
-      },
-      "outputs": [],
-      "source": [
-        "!pip install torch==\"2.1.2\"\n",
-        "!pip install -e git+https://github.com/OpenAccess-AI-Collective/axolotl#egg=axolotl\n",
-        "!pip install flash-attn==\"2.5.0\"\n",
-        "!pip install deepspeed==\"0.13.1\""
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "BW2MFr7HTjub"
-      },
-      "source": [
-        "## Create an yaml config file"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "9pkF2dSoQEUN"
-      },
-      "outputs": [],
-      "source": [
-        "import yaml\n",
-        "\n",
-        "# Your YAML string\n",
-        "yaml_string = \"\"\"\n",
-        "base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\n",
-        "model_type: LlamaForCausalLM\n",
-        "tokenizer_type: LlamaTokenizer\n",
-        "is_llama_derived_model: true\n",
-        "\n",
-        "load_in_8bit: false\n",
-        "load_in_4bit: true\n",
-        "strict: false\n",
-        "\n",
-        "datasets:\n",
-        "  - path: mhenrichsen/alpaca_2k_test\n",
-        "    type: alpaca\n",
-        "dataset_prepared_path:\n",
-        "val_set_size: 0.05\n",
-        "output_dir: ./qlora-out\n",
-        "\n",
-        "adapter: qlora\n",
-        "lora_model_dir:\n",
-        "\n",
-        "sequence_len: 1096\n",
-        "sample_packing: true\n",
-        "pad_to_sequence_len: true\n",
-        "\n",
-        "lora_r: 32\n",
-        "lora_alpha: 16\n",
-        "lora_dropout: 0.05\n",
-        "lora_target_modules:\n",
-        "lora_target_linear: true\n",
-        "lora_fan_in_fan_out:\n",
-        "\n",
-        "wandb_project:\n",
-        "wandb_entity:\n",
-        "wandb_watch:\n",
-        "wandb_name:\n",
-        "wandb_log_model:\n",
-        "\n",
-        "mlflow_experiment_name: colab-example\n",
-        "\n",
-        "gradient_accumulation_steps: 1\n",
-        "micro_batch_size: 1\n",
-        "num_epochs: 4\n",
-        "max_steps: 20\n",
-        "optimizer: paged_adamw_32bit\n",
-        "lr_scheduler: cosine\n",
-        "learning_rate: 0.0002\n",
-        "\n",
-        "train_on_inputs: false\n",
-        "group_by_length: false\n",
-        "bf16: false\n",
-        "fp16: true\n",
-        "tf32: false\n",
-        "\n",
-        "gradient_checkpointing: true\n",
-        "early_stopping_patience:\n",
-        "resume_from_checkpoint:\n",
-        "local_rank:\n",
-        "logging_steps: 1\n",
-        "xformers_attention:\n",
-        "flash_attention: false\n",
-        "\n",
-        "warmup_steps: 10\n",
-        "evals_per_epoch:\n",
-        "saves_per_epoch:\n",
-        "debug:\n",
-        "deepspeed:\n",
-        "weight_decay: 0.0\n",
-        "fsdp:\n",
-        "fsdp_config:\n",
-        "special_tokens:\n",
-        "\n",
-        "\"\"\"\n",
-        "\n",
-        "# Convert the YAML string to a Python dictionary\n",
-        "yaml_dict = yaml.safe_load(yaml_string)\n",
-        "\n",
-        "# Specify your file path\n",
-        "file_path = 'test_axolotl.yaml'\n",
-        "\n",
-        "# Write the YAML file\n",
-        "with open(file_path, 'w') as file:\n",
-        "    yaml.dump(yaml_dict, file)\n"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {
-        "id": "bidoj8YLTusD"
-      },
-      "source": [
-        "## Launch the training"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "colab": {
-          "base_uri": "https://localhost:8080/"
-        },
-        "id": "ydTI2Jk2RStU",
-        "outputId": "d6d0df17-4b53-439c-c802-22c0456d301b"
-      },
-      "outputs": [],
-      "source": [
-        "# Buy using the ! the comand will be executed as a bash command\n",
-        "!accelerate launch -m axolotl.cli.train /content/test_axolotl.yaml"
-      ]
-    },
-    {
-      "cell_type": "markdown",
-      "metadata": {},
-      "source": [
-        "## Play with inference"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {},
-      "outputs": [],
-      "source": [
-        "# Buy using the ! the comand will be executed as a bash command\n",
-        "!accelerate launch -m axolotl.cli.inference /content/test_axolotl.yaml \\\n",
-        "    --qlora_model_dir=\"./qlora-out\" --gradio"
-      ]
-    }
-  ],
-  "metadata": {
-    "accelerator": "GPU",
-    "colab": {
-      "gpuType": "T4",
-      "provenance": []
-    },
-    "kernelspec": {
-      "display_name": "Python 3",
-      "name": "python3"
-    },
-    "language_info": {
-      "name": "python"
-    }
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "AKjdG7tbTb-n"
+   },
+   "source": [
+    "# Example notebook for running Axolotl on google colab"
+   ]
   },
-  "nbformat": 4,
-  "nbformat_minor": 0
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "RcbNpOgWRcii"
+   },
+   "outputs": [],
+   "source": [
+    "import torch\n",
+    "# Check so there is a gpu available, a T4(free tier) is enough to run this notebook\n",
+    "assert (torch.cuda.is_available()==True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "h3nLav8oTRA5"
+   },
+   "source": [
+    "## Install Axolotl and dependencies"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "3c3yGAwnOIdi",
+    "outputId": "e3777b5a-40ef-424f-e181-62dfecd1dd01"
+   },
+   "outputs": [],
+   "source": [
+    "!pip install torch==\"2.1.2\"\n",
+    "!pip install -e git+https://github.com/OpenAccess-AI-Collective/axolotl#egg=axolotl\n",
+    "!pip install flash-attn==\"2.5.0\"\n",
+    "!pip install deepspeed==\"0.13.1\"!pip install mlflow==\"2.13.0\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "BW2MFr7HTjub"
+   },
+   "source": [
+    "## Create an yaml config file"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "id": "9pkF2dSoQEUN"
+   },
+   "outputs": [],
+   "source": [
+    "import yaml\n",
+    "\n",
+    "# Your YAML string\n",
+    "yaml_string = \"\"\"\n",
+    "base_model: TinyLlama/TinyLlama-1.1B-intermediate-step-1431k-3T\n",
+    "model_type: LlamaForCausalLM\n",
+    "tokenizer_type: LlamaTokenizer\n",
+    "\n",
+    "load_in_8bit: false\n",
+    "load_in_4bit: true\n",
+    "strict: false\n",
+    "\n",
+    "datasets:\n",
+    "  - path: mhenrichsen/alpaca_2k_test\n",
+    "    type: alpaca\n",
+    "dataset_prepared_path:\n",
+    "val_set_size: 0.05\n",
+    "output_dir: ./outputs/qlora-out\n",
+    "\n",
+    "adapter: qlora\n",
+    "lora_model_dir:\n",
+    "\n",
+    "sequence_len: 4096\n",
+    "sample_packing: true\n",
+    "eval_sample_packing: false\n",
+    "pad_to_sequence_len: true\n",
+    "\n",
+    "lora_r: 32\n",
+    "lora_alpha: 16\n",
+    "lora_dropout: 0.05\n",
+    "lora_target_modules:\n",
+    "lora_target_linear: true\n",
+    "lora_fan_in_fan_out:\n",
+    "\n",
+    "wandb_project:\n",
+    "wandb_entity:\n",
+    "wandb_watch:\n",
+    "wandb_name:\n",
+    "wandb_log_model:\n",
+    "\n",
+    "gradient_accumulation_steps: 4\n",
+    "micro_batch_size: 2\n",
+    "num_epochs: 4\n",
+    "optimizer: paged_adamw_32bit\n",
+    "lr_scheduler: cosine\n",
+    "learning_rate: 0.0002\n",
+    "\n",
+    "train_on_inputs: false\n",
+    "group_by_length: false\n",
+    "bf16: auto\n",
+    "fp16:\n",
+    "tf32: false\n",
+    "\n",
+    "gradient_checkpointing: true\n",
+    "early_stopping_patience:\n",
+    "resume_from_checkpoint:\n",
+    "local_rank:\n",
+    "logging_steps: 1\n",
+    "xformers_attention:\n",
+    "flash_attention: true\n",
+    "\n",
+    "warmup_steps: 10\n",
+    "evals_per_epoch: 4\n",
+    "saves_per_epoch: 1\n",
+    "debug:\n",
+    "deepspeed:\n",
+    "weight_decay: 0.0\n",
+    "fsdp:\n",
+    "fsdp_config:\n",
+    "special_tokens:\n",
+    "\n",
+    "\"\"\"\n",
+    "\n",
+    "# Convert the YAML string to a Python dictionary\n",
+    "yaml_dict = yaml.safe_load(yaml_string)\n",
+    "\n",
+    "# Specify your file path\n",
+    "file_path = 'test_axolotl.yaml'\n",
+    "\n",
+    "# Write the YAML file\n",
+    "with open(file_path, 'w') as file:\n",
+    "    yaml.dump(yaml_dict, file)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "id": "bidoj8YLTusD"
+   },
+   "source": [
+    "## Launch the training"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "colab": {
+     "base_uri": "https://localhost:8080/"
+    },
+    "id": "ydTI2Jk2RStU",
+    "outputId": "d6d0df17-4b53-439c-c802-22c0456d301b"
+   },
+   "outputs": [],
+   "source": [
+    "# Buy using the ! the comand will be executed as a bash command\n",
+    "!accelerate launch -m axolotl.cli.train /content/test_axolotl.yaml"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Play with inference"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Buy using the ! the comand will be executed as a bash command\n",
+    "!accelerate launch -m axolotl.cli.inference /content/test_axolotl.yaml \\\n",
+    "    --qlora_model_dir=\"./qlora-out\" --gradio"
+   ]
+  }
+ ],
+ "metadata": {
+  "accelerator": "GPU",
+  "colab": {
+   "gpuType": "T4",
+   "provenance": []
+  },
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.1"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
 }

examples/dbrx/16bit-lora.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 512
 sample_packing: false

examples/dbrx/8bit-lora.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 512
 sample_packing: false

examples/dbrx/fft-ds-zero3.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 512
 sample_packing: false

examples/falcon/config-7b-lora.yml

@@ -28,7 +28,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./falcon-7b
+output_dir: ./outputs/falcon-7b
 batch_size: 2
 micro_batch_size: 1
 num_epochs: 4

examples/falcon/config-7b-qlora.yml

@@ -42,7 +42,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 # QLoRA paper Table 9
 # - 16 for 7b & 13b

examples/falcon/config-7b.yml

@@ -28,7 +28,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./falcon-7b
+output_dir: ./outputs/falcon-7b
 batch_size: 2
 micro_batch_size: 1
 num_epochs: 4

examples/gemma/qlora.yml

@@ -12,7 +12,7 @@ datasets:
   - path: mhenrichsen/alpaca_2k_test
     type: alpaca
 val_set_size: 0.1
-output_dir: ./out
+output_dir: ./outputs/out
 
 adapter: qlora
 lora_r: 32

examples/gptj/qlora.yml

@@ -23,7 +23,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 gradient_accumulation_steps: 2
 micro_batch_size: 2
 num_epochs: 2

examples/jamba/qlora.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: false

examples/jamba/qlora_deepspeed.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: false

examples/jeopardy-bot/config.yml

@@ -21,7 +21,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./jeopardy-bot-7b
+output_dir: ./outputs/jeopardy-bot-7b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4

examples/llama-2/fft_optimized.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: true

examples/llama-2/gptq-lora.yml

@@ -33,7 +33,7 @@ wandb_project:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./model-out
+output_dir: ./outputs/model-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4

examples/llama-2/lisa.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./lisa-out
+output_dir: ./outputs/lisa-out
 
 sequence_len: 4096
 sample_packing: true

examples/llama-2/loftq.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/llama-2/lora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/llama-2/qlora-fsdp.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/llama-2/qlora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/llama-2/relora.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./relora-out
+output_dir: ./outputs/relora-out
 
 adapter: qlora
 lora_model_dir:

examples/llama-3/fft-8b.yaml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 8192
 sample_packing: true

examples/llama-3/instruct-lora-8b.yml

@@ -0,0 +1,76 @@
+base_model: meta-llama/Meta-Llama-3-8B-Instruct
+model_type: LlamaForCausalLM
+tokenizer_type: AutoTokenizer
+
+load_in_8bit: true
+load_in_4bit: false
+strict: false
+
+chat_template: llama3
+datasets:
+  - path: fozziethebeat/alpaca_messages_2k_test
+    type: chat_template
+    chat_template: llama3
+    field_messages: messages
+    message_field_role: role
+    message_field_content: content
+    roles:
+      user:
+        - user
+      assistant:
+        - assistant
+
+dataset_prepared_path:
+val_set_size: 0.05
+output_dir: ./outputs/lora-out
+
+sequence_len: 4096
+sample_packing: false
+pad_to_sequence_len: true
+
+adapter: lora
+lora_model_dir:
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+lora_fan_in_fan_out:
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 2
+num_epochs: 4
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: false
+
+gradient_checkpointing: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+s2_attention:
+
+warmup_steps: 10
+evals_per_epoch: 4
+eval_table_size:
+eval_max_new_tokens: 128
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:

examples/llama-3/lora-8b.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true
@@ -24,6 +24,9 @@ lora_alpha: 16
 lora_dropout: 0.05
 lora_target_linear: true
 lora_fan_in_fan_out:
+lora_modules_to_save:
+  - embed_tokens
+  - lm_head
 
 wandb_project:
 wandb_entity:

examples/llama-3/qlora-fsdp-70b.yaml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out/qlora-llama3-70b
+output_dir: ./outputs/out/qlora-llama3-70b
 
 adapter: qlora
 lora_model_dir:

examples/llama-3/qlora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/mamba/config.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 2048
 sample_packing: false

examples/mistral/bigstral-ds-zero3.yaml

@@ -23,7 +23,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 2048
 sample_packing: true

examples/mistral/config.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 8192
 sample_packing: true

examples/mistral/lora-mps.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 eval_sample_packing: false
 
 adapter: lora

examples/mistral/lora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 adapter: lora
 lora_model_dir:

examples/mistral/mistral-qlora-fsdp.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.02
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 model_config:
   output_router_logits: true

examples/mistral/mistral-qlora-orpo.yml

@@ -16,7 +16,7 @@ datasets:
     type: chat_template.argilla
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
-output_dir: ./mistral-qlora-orpo-out
+output_dir: ./outputs/mistral-qlora-orpo-out
 
 adapter: qlora
 lora_model_dir:

examples/mistral/mixtral-8x22b-qlora-fsdp.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.02
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 model_config:
   output_router_logits: true

examples/mistral/mixtral-qlora-fsdp.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.02
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 model_config:
   output_router_logits: true

examples/mistral/mixtral.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.0
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 ## You can optionally freeze the entire model and unfreeze a subset of parameters
 unfrozen_parameters:

examples/mistral/mixtral_22.yml

@@ -21,7 +21,7 @@ model_config:
 datasets:
   - path: yahma/alpaca-cleaned
     type: alpaca
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 8000
 sample_packing: true

examples/mistral/qlora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.1
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:

examples/mpt-7b/config.yml

@@ -23,7 +23,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./mpt-alpaca-7b
+output_dir: ./outputs/mpt-alpaca-7b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4

examples/openllama-3b/config.yml

@@ -25,7 +25,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./openllama-out
+output_dir: ./outputs/openllama-out
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 4

examples/openllama-3b/lora.yml

@@ -31,7 +31,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4

examples/openllama-3b/qlora.yml

@@ -25,7 +25,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 gradient_accumulation_steps: 1
 micro_batch_size: 2
 num_epochs: 4

examples/phi/phi-ft.yml

@@ -12,7 +12,7 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./phi-sft-out
+output_dir: ./outputs/phi-sft-out
 
 sequence_len: 2048
 sample_packing: true

examples/phi/phi-qlora.yml

@@ -12,7 +12,7 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./phi-sft-out
+output_dir: ./outputs/phi-sft-out
 
 sequence_len: 2048
 sample_packing: true

examples/phi/phi2-ft.yml

@@ -12,7 +12,7 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./phi-sft-out
+output_dir: ./outputs/phi-sft-out
 
 sequence_len: 2048
 sample_packing: true

examples/pythia-12b/config.yml

@@ -26,7 +26,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./pythia-12b
+output_dir: ./outputs/pythia-12b
 gradient_accumulation_steps: 1
 micro_batch_size: 1
 num_epochs: 5

examples/pythia/lora.yml

@@ -20,7 +20,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./lora-alpaca-pythia
+output_dir: ./outputs/lora-alpaca-pythia
 gradient_accumulation_steps: 1
 micro_batch_size: 4
 num_epochs: 4

examples/qwen/lora.yml

@@ -13,7 +13,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 2048  # supports up to 8192
 sample_packing: false

examples/qwen/qlora.yml

@@ -13,7 +13,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 2048  # supports up to 8192
 sample_packing: false

examples/qwen/qwen2-moe-lora.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 1024  # supports up to 32k
 sample_packing: false

examples/qwen/qwen2-moe-qlora.yaml

@@ -10,7 +10,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 1024  # supports up to 32k
 sample_packing: false

examples/redpajama/config-3b.yml

@@ -24,7 +24,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./redpajama-alpaca-3b
+output_dir: ./outputs/redpajama-alpaca-3b
 batch_size: 4
 micro_batch_size: 1
 num_epochs: 4

examples/replit-3b/config-lora.yml

@@ -23,7 +23,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./lora-replit
+output_dir: ./outputs/lora-replit
 batch_size: 8
 micro_batch_size: 1
 num_epochs: 4

examples/stablelm-2/1.6b/fft.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path: last_run_prepared
 val_set_size: 0.05
-output_dir: ./out
+output_dir: ./outputs/out
 
 sequence_len: 4096
 sample_packing: true

examples/stablelm-2/1.6b/lora.yml

@@ -12,7 +12,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/starcoder2/qlora.yml

@@ -11,7 +11,7 @@ datasets:
 
 dataset_prepared_path:
 val_set_size: 0.2
-output_dir: ./qlora
+output_dir: ./outputs/qlora
 
 adapter: qlora
 lora_model_dir:

examples/tiny-llama/lora-mps.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/tiny-llama/lora.yml

@@ -11,7 +11,7 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./lora-out
+output_dir: ./outputs/lora-out
 
 sequence_len: 4096
 sample_packing: true

examples/tiny-llama/pretrain.yml

@@ -14,7 +14,7 @@ pretraining_dataset:
   type: pretrain
 dataset_prepared_path:
 val_set_size: 0.0
-output_dir: ./model-out
+output_dir: ./outputs/model-out
 
 sequence_len: 2048
 sample_packing: true

examples/tiny-llama/qlora.yml

@@ -11,13 +11,14 @@ datasets:
     type: alpaca
 dataset_prepared_path:
 val_set_size: 0.05
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 adapter: qlora
 lora_model_dir:
 
 sequence_len: 4096
 sample_packing: true
+eval_sample_packing: false
 pad_to_sequence_len: true
 
 lora_r: 32

examples/xgen-7b/xgen-7b-8k-qlora.yml

@@ -40,7 +40,7 @@ wandb_entity:
 wandb_watch:
 wandb_name:
 wandb_log_model:
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 
 # QLoRA paper Table 9
 # - 16 for 7b & 13b

examples/yi-34B-chat/qlora.yml

@@ -33,7 +33,7 @@ eval_sample_packing: false
 eval_batch_size: 1
 
 # LoRA
-output_dir: ./qlora-out
+output_dir: ./outputs/qlora-out
 adapter: qlora
 lora_model_dir:
 lora_r: 32

requirements.txt

@@ -1,22 +1,22 @@
 --extra-index-url https://huggingface.github.io/autogptq-index/whl/cu118/
 packaging==23.2
-peft==0.10.0
-transformers @ git+https://github.com/huggingface/transformers.git@43d17c18360ac9c3d3491389328e2fe55fe8f9ce
-tokenizers==0.15.0
-bitsandbytes==0.43.0
-accelerate==0.28.0
-deepspeed==0.13.1
+peft==0.11.1
+transformers==4.41.1
+tokenizers==0.19.1
+bitsandbytes==0.43.1
+accelerate==0.30.1
+deepspeed==0.14.2
 pydantic==2.6.3
 addict
 fire
 PyYAML>=6.0
 requests
-datasets==2.15.0
-flash-attn==2.5.5
+datasets==2.19.1
+flash-attn==2.5.8
 sentencepiece
 wandb
 einops
-xformers==0.0.22
+xformers==0.0.26.post1
 optimum==1.16.2
 hf_transfer
 colorama
@@ -28,7 +28,7 @@ scipy
 scikit-learn==1.2.2
 pynvml
 art
-fschat @ git+https://github.com/lm-sys/FastChat.git@5095615810cf613dba7f27dd155f571fcff976d8
+fschat @ git+https://github.com/lm-sys/FastChat.git@27a05b04a35510afb1d767ae7e5990cbd278f8fe
 gradio==3.50.2
 tensorboard
 
@@ -39,6 +39,6 @@ s3fs
 gcsfs
 # adlfs
 
-trl==0.8.5
+trl==0.8.6
 zstandard==0.22.0
 fastcore

scripts/cloud-entrypoint-term.sh

@@ -0,0 +1,82 @@
+#!/bin/bash
+
+# Export specific ENV variables to /etc/rp_environment
+echo "Exporting environment variables..."
+printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
+conda init
+# this needs to come after conda init
+echo 'source /etc/rp_environment' >> ~/.bashrc
+
+add_keys_to_authorized() {
+    local key_value=$1
+
+    # Create the ~/.ssh directory and set permissions
+    mkdir -p ~/.ssh
+    chmod 700 ~/.ssh
+
+    # Create the authorized_keys file if it doesn't exist
+    touch ~/.ssh/authorized_keys
+
+    # Initialize an empty key variable
+    local key=""
+
+    # Read the key variable word by word
+    for word in $key_value; do
+        # Check if the word looks like the start of a key
+        if [[ $word == ssh-* ]]; then
+            # If there's a key being built, add it to the authorized_keys file
+            if [[ -n $key ]]; then
+                echo $key >> ~/.ssh/authorized_keys
+            fi
+            # Start a new key
+            key=$word
+        else
+            # Append the word to the current key
+            key="$key $word"
+        fi
+    done
+
+    # Add the last key to the authorized_keys file
+    if [[ -n $key ]]; then
+        echo $key >> ~/.ssh/authorized_keys
+    fi
+
+    # Set the correct permissions
+    chmod 600 ~/.ssh/authorized_keys
+    chmod 700 -R ~/.ssh
+}
+
+if [[ $PUBLIC_KEY ]]; then
+    # runpod
+    add_keys_to_authorized "$PUBLIC_KEY"
+    # Start the SSH service in the background
+    service ssh start
+elif [[ $SSH_KEY ]]; then
+    # latitude.sh
+    add_keys_to_authorized "$SSH_KEY"
+    # Start the SSH service in the background
+    service ssh start
+else
+    echo "No PUBLIC_KEY or SSH_KEY environment variable provided, not starting openSSH daemon"
+fi
+
+# Check if JUPYTER_PASSWORD is set and not empty
+if [ -n "$JUPYTER_PASSWORD" ]; then
+    # Set JUPYTER_TOKEN to the value of JUPYTER_PASSWORD
+    export JUPYTER_TOKEN="$JUPYTER_PASSWORD"
+fi
+
+if [ "$JUPYTER_DISABLE" != "1" ]; then
+    # Run Jupyter Lab in the background
+    jupyter lab --port=8888 --ip=* --allow-root --ServerApp.allow_origin=* &
+fi
+
+if [ ! -d "/workspace/data/axolotl-artifacts" ]; then
+    mkdir -p /workspace/data/axolotl-artifacts
+fi
+if [ ! -L "/workspace/axolotl/outputs" ]; then
+    ln -sf /workspace/data/axolotl-artifacts /workspace/axolotl/outputs
+fi
+
+# Execute the passed arguments (CMD)
+exec "$@"

scripts/cloud-entrypoint.sh

@@ -5,20 +5,53 @@ echo "Exporting environment variables..."
 printenv | grep -E '^RUNPOD_|^PATH=|^_=' | sed 's/^\(.*\)=\(.*\)$/export \1="\2"/' >> /etc/rp_environment
 echo 'source /etc/rp_environment' >> ~/.bashrc
 
+add_keys_to_authorized() {
+    local key_value=$1
+
+    # Create the ~/.ssh directory and set permissions
+    mkdir -p ~/.ssh
+    chmod 700 ~/.ssh
+
+    # Create the authorized_keys file if it doesn't exist
+    touch ~/.ssh/authorized_keys
+
+    # Initialize an empty key variable
+    local key=""
+
+    # Read the key variable word by word
+    for word in $key_value; do
+        # Check if the word looks like the start of a key
+        if [[ $word == ssh-* ]]; then
+            # If there's a key being built, add it to the authorized_keys file
+            if [[ -n $key ]]; then
+                echo $key >> ~/.ssh/authorized_keys
+            fi
+            # Start a new key
+            key=$word
+        else
+            # Append the word to the current key
+            key="$key $word"
+        fi
+    done
+
+    # Add the last key to the authorized_keys file
+    if [[ -n $key ]]; then
+        echo $key >> ~/.ssh/authorized_keys
+    fi
+
+    # Set the correct permissions
+    chmod 600 ~/.ssh/authorized_keys
+    chmod 700 -R ~/.ssh
+}
+
 if [[ $PUBLIC_KEY ]]; then
     # runpod
-    mkdir -p ~/.ssh
-    chmod 700 ~/.ssh
-    echo $PUBLIC_KEY >> ~/.ssh/authorized_keys
-    chmod 700 -R ~/.ssh
+    add_keys_to_authorized "$PUBLIC_KEY"
     # Start the SSH service in the background
     service ssh start
-elif [ -n "$SSH_KEY" ]; then
+elif [[ $SSH_KEY ]]; then
     # latitude.sh
-    mkdir -p ~/.ssh
-    chmod 700 ~/.ssh
-    echo $SSH_KEY >> ~/.ssh/authorized_keys
-    chmod 700 -R ~/.ssh
+    add_keys_to_authorized "$SSH_KEY"
     # Start the SSH service in the background
     service ssh start
 else
@@ -36,5 +69,12 @@ if [ "$JUPYTER_DISABLE" != "1" ]; then
     jupyter lab --port=8888 --ip=* --allow-root --ServerApp.allow_origin=* &
 fi
 
+if [ ! -d "/workspace/data/axolotl-artifacts" ]; then
+    mkdir -p /workspace/data/axolotl-artifacts
+fi
+if [ ! -L "/workspace/axolotl/outputs" ]; then
+    ln -sf /workspace/data/axolotl-artifacts /workspace/axolotl/outputs
+fi
+
 # Execute the passed arguments (CMD)
 exec "$@"

setup.py

@@ -30,8 +30,11 @@ def parse_requirements():
 
     try:
         if "Darwin" in platform.system():
-            _install_requires.pop(_install_requires.index("xformers==0.0.22"))
+            # don't install xformers on MacOS
+            _install_requires.pop(_install_requires.index("xformers==0.0.26.post1"))
         else:
+            # detect the version of torch already installed
+            # and set it so dependencies don't clobber the torch version
             torch_version = version("torch")
             _install_requires.append(f"torch=={torch_version}")
 
@@ -45,9 +48,15 @@ def parse_requirements():
             else:
                 raise ValueError("Invalid version format")
 
-            if (major, minor) >= (2, 1):
-                _install_requires.pop(_install_requires.index("xformers==0.0.22"))
-                _install_requires.append("xformers>=0.0.23")
+            if (major, minor) >= (2, 3):
+                pass
+            elif (major, minor) >= (2, 2):
+                _install_requires.pop(_install_requires.index("xformers==0.0.26.post1"))
+                _install_requires.append("xformers>=0.0.25.post1")
+            else:
+                _install_requires.pop(_install_requires.index("xformers==0.0.26.post1"))
+                _install_requires.append("xformers>=0.0.23.post1")
+
     except PackageNotFoundError:
         pass
 
@@ -59,7 +68,7 @@ install_requires, dependency_links = parse_requirements()
 
 setup(
     name="axolotl",
-    version="0.4.0",
+    version="0.4.1",
     description="LLM Trainer",
     long_description="Axolotl is a tool designed to streamline the fine-tuning of various AI models, offering support for multiple configurations and architectures.",
     package_dir={"": "src"},
@@ -68,13 +77,13 @@ setup(
     dependency_links=dependency_links,
     extras_require={
         "flash-attn": [
-            "flash-attn==2.5.5",
+            "flash-attn==2.5.8",
         ],
         "fused-dense-lib": [
-            "fused-dense-lib  @ git+https://github.com/Dao-AILab/flash-attention@v2.3.3#subdirectory=csrc/fused_dense_lib",
+            "fused-dense-lib  @ git+https://github.com/Dao-AILab/flash-attention@v2.5.8#subdirectory=csrc/fused_dense_lib",
         ],
         "deepspeed": [
-            "deepspeed==0.13.1",
+            "deepspeed==0.14.2",
             "deepspeed-kernels",
         ],
         "mamba-ssm": [

src/axolotl/cli/merge_lora.py

@@ -25,6 +25,8 @@ def do_cli(config: Path = Path("examples/"), **kwargs):
         load_in_8bit=False,
         load_in_4bit=False,
         flash_attention=False,
+        deepspeed=None,
+        fsdp=None,
         **kwargs,
     )
 
@@ -40,6 +42,7 @@ def do_cli(config: Path = Path("examples/"), **kwargs):
     parsed_cfg.flash_attention = False
     parsed_cfg.deepspeed = None
     parsed_cfg.fsdp = None
+    parsed_cfg.fsdp_config = None
 
     do_merge_lora(cfg=parsed_cfg, cli_args=parsed_cli_args)
 

src/axolotl/cli/preprocess.py

@@ -19,7 +19,10 @@ from axolotl.cli import (
 )
 from axolotl.common.cli import PreprocessCliArgs
 from axolotl.common.const import DEFAULT_DATASET_PREPARED_PATH
-from axolotl.prompt_strategies.sharegpt import register_chatml_template
+from axolotl.prompt_strategies.sharegpt import (
+    register_chatml_template,
+    register_llama3_template,
+)
 
 LOG = logging.getLogger("axolotl.cli.preprocess")
 
@@ -36,13 +39,22 @@ def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs):
         return_remaining_strings=True
     )
 
-    if parsed_cfg.chat_template == "chatml" and parsed_cfg.default_system_message:
-        LOG.info(
-            f"ChatML set. Adding default system message: {parsed_cfg.default_system_message}"
-        )
-        register_chatml_template(parsed_cfg.default_system_message)
-    else:
-        register_chatml_template()
+    if parsed_cfg.chat_template == "chatml":
+        if parsed_cfg.default_system_message:
+            LOG.info(
+                f"ChatML set. Adding default system message: {parsed_cfg.default_system_message}"
+            )
+            register_chatml_template(parsed_cfg.default_system_message)
+        else:
+            register_chatml_template()
+    elif parsed_cfg.chat_template == "llama3":
+        if parsed_cfg.default_system_message:
+            LOG.info(
+                f"LLaMA-3 set. Adding default system message: {parsed_cfg.default_system_message}"
+            )
+            register_llama3_template(parsed_cfg.default_system_message)
+        else:
+            register_llama3_template()
 
     if not parsed_cfg.dataset_prepared_path:
         msg = (

src/axolotl/cli/train.py

@@ -19,7 +19,10 @@ from axolotl.cli import (
     print_axolotl_text_art,
 )
 from axolotl.common.cli import TrainerCliArgs
-from axolotl.prompt_strategies.sharegpt import register_chatml_template
+from axolotl.prompt_strategies.sharegpt import (
+    register_chatml_template,
+    register_llama3_template,
+)
 from axolotl.train import train
 
 LOG = logging.getLogger("axolotl.cli.train")
@@ -47,6 +50,14 @@ def do_train(cfg, cli_args) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
     else:
         register_chatml_template()
 
+    if cfg.chat_template == "llama3" and cfg.default_system_message:
+        LOG.info(
+            f"LLaMA-3 set. Adding default system message: {cfg.default_system_message}"
+        )
+        register_llama3_template(cfg.default_system_message)
+    else:
+        register_llama3_template()
+
     if cfg.rl:  # and cfg.rl != "orpo":
         dataset_meta = load_rl_datasets(cfg=cfg, cli_args=cli_args)
     else:

src/axolotl/core/trainer_builder.py

@@ -30,7 +30,7 @@ from transformers import (
 )
 from transformers.trainer_utils import seed_worker
 from transformers.utils import is_sagemaker_mp_enabled
-from trl import DPOTrainer, ORPOConfig, ORPOTrainer
+from trl import DPOTrainer, KTOConfig, KTOTrainer, ORPOConfig, ORPOTrainer
 from trl.trainer.utils import pad_to_length
 
 from axolotl.loraplus import create_loraplus_optimizer
@@ -43,7 +43,7 @@ from axolotl.utils.callbacks import (
     LossWatchDogCallback,
     SaveAxolotlConfigtoWandBCallback,
     SaveBetterTransformerModelCallback,
-    SaveModelOnTrainEndCallback,
+    SaveModelCallback,
     bench_eval_callback_factory,
     causal_lm_bench_eval_callback_factory,
     log_prediction_callback_factory,
@@ -91,11 +91,12 @@ def _sanitize_kwargs_for_tagging(tag_names, kwargs=None):
 
 
 @dataclass
-class AxolotlTrainingArguments(TrainingArguments):
+class AxolotlTrainingMixins:
     """
-    Extend the base TrainingArguments for axolotl helpers
+    Mixin class for the Axolotl training args.
     """
 
+    # pylint: disable=duplicate-code
     model_type: Optional[str] = field(
         default=None, metadata={"help": "HF model configuration model_type."}
     )
@@ -125,14 +126,22 @@ class AxolotlTrainingArguments(TrainingArguments):
         default=1.0,
         metadata={"help": "Sample packing efficiency for calculating batch length."},
     )
+    sample_packing_bin_size: int = field(
+        default=200,
+        metadata={
+            "help": "The max number of samples that packed sample can contain after packing. Increase for better packing."
+        },
+    )
+    sample_packing_group_size: int = field(
+        default=100000,
+        metadata={
+            "help": "The number of samples to group together for packing. Increase for better packing."
+        },
+    )
     max_seq_length: int = field(
         default=2048,
         metadata={"help": "The maximum sequence length the model can handle"},
     )
-    sample_packing_seq_len_multiplier: int = field(
-        default=1,
-        metadata={"help": "the multiplier for the max len for packed sequences"},
-    )
     relora_steps: Optional[int] = field(
         default=None,
         metadata={"help": "how often to reset for ReLoRA"},
@@ -219,6 +228,30 @@ class AxolotlTrainingArguments(TrainingArguments):
     )
 
 
+@dataclass
+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
+    so it can't be used as a mixin.
+    """
+
+
+@dataclass
+class AxolotlORPOConfig(AxolotlTrainingMixins, ORPOConfig):
+    """
+    ORPO config for ORPO training
+    """
+
+
+@dataclass
+class AxolotlKTOConfig(AxolotlTrainingMixins, KTOConfig):
+    """
+    KTO config for KTO training
+    """
+
+
 class AxolotlTrainer(Trainer):
     """
     Extend the base Trainer for axolotl helpers
@@ -346,11 +379,12 @@ class AxolotlTrainer(Trainer):
                 )
             return MultipackBatchSampler(
                 RandomSampler(self.train_dataset),
-                batch_size=batch_size,
-                drop_last=True,
-                batch_max_len=batch_max_len,
                 lengths=get_dataset_lengths(self.train_dataset),
-                packing_efficiency_estimate=self.args.sample_packing_efficiency,
+                batch_max_len=batch_max_len,
+                batch_size=batch_size,
+                group_size=self.args.sample_packing_group_size,
+                bin_size=self.args.sample_packing_bin_size,
+                drop_last=True,
             )
         if self.args.curriculum_sampling:
             return SequentialSampler(self.train_dataset)
@@ -370,11 +404,12 @@ class AxolotlTrainer(Trainer):
                 )
             return MultipackBatchSampler(
                 SequentialSampler(eval_dataset),
-                batch_size=batch_size,
-                drop_last=True,
+                lengths=get_dataset_lengths(self.eval_dataset),
                 batch_max_len=batch_max_len,
-                lengths=get_dataset_lengths(eval_dataset),
-                packing_efficiency_estimate=self.args.sample_packing_efficiency,
+                batch_size=batch_size,
+                group_size=self.args.sample_packing_group_size,
+                bin_size=self.args.sample_packing_bin_size,
+                drop_last=True,
             )
         return super()._get_eval_sampler(eval_dataset)
 
@@ -798,6 +833,40 @@ class AxolotlDPOTrainer(DPOTrainer):
 
     tag_names = ["axolotl", "dpo"]
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.optimizer = None
+
+    def create_optimizer(self):
+        if self.args.loraplus_lr_ratio is None:
+            return super().create_optimizer()
+
+        opt_model = self.model_wrapped if is_sagemaker_mp_enabled() else self.model
+        if self.optimizer is None:  # pylint: disable=access-member-before-definition
+            optimizer_cls, optimizer_kwargs = Trainer.get_optimizer_cls_and_kwargs(
+                self.args,
+                opt_model,
+            )
+
+            loraplus_lr_ratio = getattr(self.args, "loraplus_lr_ratio", None)
+            if loraplus_lr_ratio:
+                print("Using lora+")
+            loraplus_lr_embedding = getattr(self.args, "loraplus_lr_embedding", None)
+            self.optimizer = create_loraplus_optimizer(  # pylint: disable=attribute-defined-outside-init
+                opt_model,
+                optimizer_cls,
+                optimizer_kwargs,
+                loraplus_lr_ratio,
+                loraplus_lr_embedding,
+            )
+
+        if is_sagemaker_mp_enabled():
+            self.optimizer = smp.DistributedOptimizer(  # pylint: disable=attribute-defined-outside-init
+                self.optimizer
+            )
+
+        return self.optimizer
+
     @wraps(DPOTrainer.push_to_hub)
     def push_to_hub(self, *args, **kwargs) -> str:
         """
@@ -826,6 +895,14 @@ class AxolotlORPOTrainer(ORPOTrainer):
     tag_names = ["axolotl", "orpo"]
 
 
+class AxolotlKTOTrainer(KTOTrainer):
+    """
+    Extend the base KTOTrainer for axolotl helpers
+    """
+
+    tag_names = ["axolotl", "kto"]
+
+
 class TrainerBuilderBase(abc.ABC):
     """
     Base class for trainer builder
@@ -945,7 +1022,7 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         if self.cfg.loss_watchdog_threshold is not None:
             callbacks.append(LossWatchDogCallback(self.cfg))
 
-        callbacks.append(SaveModelOnTrainEndCallback())
+        callbacks.append(SaveModelCallback())
 
         return callbacks
 
@@ -1071,11 +1148,6 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         if self.cfg.save_safetensors is not None:
             training_arguments_kwargs["save_safetensors"] = self.cfg.save_safetensors
 
-        if self.cfg.sample_packing_eff_est:
-            training_arguments_kwargs[
-                "sample_packing_efficiency"
-            ] = self.cfg.sample_packing_eff_est
-
         if self.cfg.dataloader_pin_memory is not None:
             training_arguments_kwargs[
                 "dataloader_pin_memory"
@@ -1123,6 +1195,8 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
             # default to saving each epoch if not defined
             training_arguments_kwargs["save_strategy"] = "epoch"
 
+        training_arguments_kwargs["save_only_model"] = self.cfg.save_only_model
+
         if self.cfg.do_bench_eval:
             training_arguments_kwargs["do_bench_eval"] = self.cfg.do_bench_eval
             if self.cfg.bench_dataset:
@@ -1202,11 +1276,14 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         )
         training_arguments_kwargs["group_by_length"] = self.cfg.group_by_length
         training_arguments_kwargs["curriculum_sampling"] = self.cfg.curriculum_sampling
-        report_to = None
+        report_to = []
         if self.cfg.use_wandb:
-            report_to = "wandb"
+            report_to.append("wandb")
         if self.cfg.use_mlflow:
-            report_to = "mlflow"
+            report_to.append("mlflow")
+        if self.cfg.use_tensorboard:
+            report_to.append("tensorboard")
+
         training_arguments_kwargs["report_to"] = report_to
         training_arguments_kwargs["run_name"] = (
             self.cfg.wandb_name if self.cfg.use_wandb else None
@@ -1246,20 +1323,27 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         training_arguments_kwargs["weight_decay"] = (
             self.cfg.weight_decay if self.cfg.weight_decay is not None else 0.0
         )
-        training_arguments_kwargs["sample_packing"] = (
-            self.cfg.sample_packing if self.cfg.sample_packing else False
-        )
-        training_arguments_kwargs["multipack_real_batches"] = (
-            self.cfg.flash_attention is not True
-        )
-        training_arguments_kwargs["eval_sample_packing"] = (
-            self.cfg.sample_packing
-            if self.cfg.eval_sample_packing is not False
-            else False
-        )
+
+        training_arguments_kwargs["sample_packing"] = bool(self.cfg.sample_packing)
         training_arguments_kwargs[
-            "sample_packing_seq_len_multiplier"
-        ] = self.cfg.micro_batch_size
+            "multipack_real_batches"
+        ] = not self.cfg.flash_attention
+        training_arguments_kwargs["eval_sample_packing"] = bool(
+            self.cfg.eval_sample_packing
+        )
+        if self.cfg.sample_packing_bin_size is not None:
+            training_arguments_kwargs[
+                "sample_packing_bin_size"
+            ] = self.cfg.sample_packing_bin_size
+        if self.cfg.sample_packing_group_size is not None:
+            training_arguments_kwargs[
+                "sample_packing_group_size"
+            ] = self.cfg.sample_packing_group_size
+        if self.cfg.sample_packing_eff_est:
+            training_arguments_kwargs[
+                "sample_packing_efficiency"
+            ] = self.cfg.sample_packing_eff_est
+
         if self.cfg.relora_steps:
             training_arguments_kwargs["relora_steps"] = self.cfg.relora_steps
             training_arguments_kwargs[
@@ -1429,7 +1513,7 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
 
     def get_callbacks(self):
         callbacks = super().get_callbacks()
-        callbacks.append(SaveModelOnTrainEndCallback())
+        callbacks.append(SaveModelCallback())
 
         return callbacks
 
@@ -1470,6 +1554,8 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
         if self.cfg.bf16 or self.cfg.bfloat16:
             training_args_kwargs["bf16"] = True
 
+        training_args_kwargs["loraplus_lr_ratio"] = self.cfg.loraplus_lr_ratio
+        training_args_kwargs["loraplus_lr_embedding"] = self.cfg.loraplus_lr_embedding
         training_args_kwargs["lr_scheduler_type"] = (
             self.cfg.lr_scheduler if self.cfg.lr_scheduler else "cosine"
         )
@@ -1522,17 +1608,36 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
             # trl does some odd mapping of alpha to beta to reuse the beta parameter ???
             training_args_kwargs["beta"] = self.cfg.orpo_alpha
 
-        training_args_cls = TrainingArguments
+        training_args_cls = AxolotlTrainingArguments
         if self.cfg.rl == "orpo":
-            training_args_cls = ORPOConfig
+            training_args_cls = AxolotlORPOConfig
             training_args_kwargs["dataset_num_proc"] = self.cfg.dataset_processes
+            training_args_kwargs["max_length"] = self.cfg.sequence_len
+            if self.cfg.max_prompt_len:
+                training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
 
-        training_args = training_args_cls(
+        if self.cfg.rl == "kto":
+            training_args_cls = AxolotlKTOConfig
+
+            training_args_kwargs["beta"] = self.cfg.rl_beta or 0.1
+            training_args_kwargs["desirable_weight"] = (
+                self.cfg.kto_desirable_weight or 1.0
+            )
+            training_args_kwargs["undesirable_weight"] = (
+                self.cfg.kto_undesirable_weight or 1.0
+            )
+
+            training_args_kwargs["dataset_num_proc"] = self.cfg.dataset_processes
+            training_args_kwargs["max_length"] = self.cfg.sequence_len
+            if self.cfg.max_prompt_len:
+                training_args_kwargs["max_prompt_length"] = self.cfg.max_prompt_len
+
+        training_args = training_args_cls(  # pylint: disable=unexpected-keyword-arg
+            output_dir=self.cfg.output_dir,
             per_device_train_batch_size=self.cfg.micro_batch_size,
             max_steps=self.cfg.max_steps or total_num_steps,
             gradient_accumulation_steps=self.cfg.gradient_accumulation_steps,
             learning_rate=self.cfg.learning_rate,
-            output_dir=self.cfg.output_dir,
             warmup_steps=self.cfg.warmup_steps,
             logging_first_step=True,
             logging_steps=1,
@@ -1562,7 +1667,7 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
             ] = self.cfg.precompute_ref_log_probs
         if self.cfg.rl in ["dpo", "ipo", "kto_pair"]:
             trainer_cls = AxolotlDPOTrainer
-            dpo_trainer_kwargs["beta"] = self.cfg.dpo_beta or 0.1
+            dpo_trainer_kwargs["beta"] = self.cfg.rl_beta or 0.1
             trainer_cls_args = [self.model, self.model_ref]
 
             # these aren't used for the ORPO trainer
@@ -1575,6 +1680,9 @@ class HFRLTrainerBuilder(TrainerBuilderBase):
         elif self.cfg.rl == "orpo":
             trainer_cls = AxolotlORPOTrainer
             trainer_cls_args = [self.model]
+        elif self.cfg.rl == "kto":
+            trainer_cls = AxolotlKTOTrainer
+            trainer_cls_args = [self.model]
         else:
             raise ValueError(f"Unsupported RL: {self.cfg.rl}")
         dpo_trainer = trainer_cls(

src/axolotl/monkeypatch/fastchat_conversation_turns.py

@@ -123,6 +123,17 @@ def get_turns(  # pylint: disable=too-many-return-statements
             else:
                 yield role, ""
         return
+    if self.sep_style == SeparatorStyle.LLAMA3:
+        if self.system_message:
+            # For llama3, the system message is NOT incorporated into the first human instruction
+            # All messages follow <|start_header_id|>' + role + '<|end_header_id|>\n\n'+ message + '<|eot_id|>
+            yield "", system_prompt
+        for i, (role, message) in enumerate(self.messages):
+            if message:
+                yield f"<|start_header_id|>{role}<|end_header_id|>\n\n", f"{message.strip()}<|eot_id|>"
+            else:
+                yield f"<|start_header_id|>{role}<|end_header_id|>\n\n", ""
+        return
     if self.sep_style == SeparatorStyle.GEMMA:
         if self.system_message:
             raise ValueError("Gemma chat template does not support system messages")

src/axolotl/monkeypatch/mixtral/__init__.py

@@ -42,9 +42,9 @@ def patch_mixtral_moe_forward_zero3() -> None:
         return final_hidden_states, router_logits
 
     from transformers.models.mixtral.modeling_mixtral import (
-        MixtralBLockSparseTop2MLP,
+        MixtralBlockSparseTop2MLP,
         MixtralSparseMoeBlock,
     )
 
-    MixtralBLockSparseTop2MLP.forward = mlp_forward
+    MixtralBlockSparseTop2MLP.forward = mlp_forward
     MixtralSparseMoeBlock.forward = moe_forward

src/axolotl/monkeypatch/unsloth_.py

@@ -0,0 +1,267 @@
+"""module for patching with unsloth optimizations"""
+
+import inspect
+import logging
+import re
+import types
+from typing import Tuple
+
+from peft import PeftModelForCausalLM
+from transformers.models.llama.modeling_llama import (
+    LlamaFlashAttention2,
+    LlamaForCausalLM,
+)
+
+LOG = logging.getLogger("axolotl.monkeypatch.unsloth")
+
+ORIGINAL_CEL_CODE = """    if labels is not None:
+        # Shift so that tokens < n predict n
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        # Flatten the tokens
+        loss_fct = CrossEntropyLoss()
+        shift_logits = shift_logits.view(-1, self.config.vocab_size)
+        shift_labels = shift_labels.view(-1)
+        # Enable model parallelism
+        shift_labels = shift_labels.to(shift_logits.device)
+        loss = loss_fct(shift_logits, shift_labels)
+"""
+
+PATCHED_CEL_CODE = """    if labels is not None:
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        loss = fast_cross_entropy_loss(
+            logits = shift_logits,
+            labels = shift_labels,
+        )
+"""
+
+ORIGINAL_QKV_CODE = """
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+""".lstrip(
+    "\n"
+)
+
+PATCHED_QKV_CODE = """
+    query_states, key_states, value_states = self.apply_qkv(self, hidden_states)
+""".lstrip(
+    "\n"
+)
+
+ORIGINAL_O_CODE = """
+    attn_output = self.o_proj(attn_output)
+""".lstrip(
+    "\n"
+)
+
+PATCHED_O_CODE = """
+    attn_output = self.apply_o(self, attn_output)
+""".lstrip(
+    "\n"
+)
+
+
+def original_apply_qkv(self, hidden_states):
+    query_states = self.q_proj(hidden_states)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+    return query_states, key_states, value_states
+
+
+def original_apply_o(self, hidden_states):
+    attn_output = self.o_proj(hidden_states)
+    return attn_output
+
+
+def get_forward_code() -> str:
+    forward = inspect.getsource(LlamaForCausalLM.forward)
+    return forward
+
+
+def test_cel_is_patchable() -> bool:
+    forward = get_forward_code()
+    return ORIGINAL_CEL_CODE in forward
+
+
+def get_self_attn_code() -> str:
+    forward = inspect.getsource(LlamaFlashAttention2.forward)
+    return forward
+
+
+def test_self_attn_is_patchable() -> bool:
+    qkv = get_self_attn_code()
+    return ORIGINAL_QKV_CODE in qkv and ORIGINAL_QKV_CODE in qkv
+
+
+def integrate_cross_entropy_loss_patch():
+    forward = get_forward_code()
+    LlamaForCausalLM._original_forward = forward  # pylint: disable=protected-access
+    forward, _ = detab_code(forward)
+    assert ORIGINAL_CEL_CODE in forward, "Original forward code not found"
+
+    forward = forward.replace(
+        "@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)", ""
+    )
+    forward = forward.replace(
+        "@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)",
+        "",
+    )
+    forward = forward.replace(ORIGINAL_CEL_CODE, PATCHED_CEL_CODE)
+    forward = forward.replace(
+        "def forward(",
+        "def fast_cross_entropy_loss_forward(",
+        1,
+    )
+
+    # load imports necessary
+    import transformers.models.llama.modeling_llama
+
+    items_to_import = []
+    for item in dir(transformers.models.llama.modeling_llama):
+        if item in forward:
+            items_to_import.append(item)
+
+    exec(  # pylint: disable=exec-used  # nosec B102
+        "from unsloth.kernels.cross_entropy_loss import fast_cross_entropy_loss",
+        globals(),
+    )
+
+    exec(  # pylint: disable=exec-used  # nosec B102
+        "from transformers.models.llama.modeling_llama import ("
+        + ", ".join(x for x in items_to_import)
+        + ")",
+        globals(),
+    )
+    exec(forward, globals())  # pylint: disable=exec-used  # nosec B102
+    print("patching unsloth fast_cross_entropy_loss")
+    LlamaForCausalLM.forward = fast_cross_entropy_loss_forward  # pylint: disable=undefined-variable  # noqa: F821
+
+
+def detab_code(code: str) -> Tuple[str, str]:
+    spaces = re.match(r"([\s\t]{1,})", code).group(0)
+    code = re.sub(r"^" + spaces, "", code, flags=re.MULTILINE)
+    return code, spaces
+
+
+def patch_self_attn_lora():
+    self_attn_forward = get_self_attn_code()
+    LlamaFlashAttention2._original_forward = (  # pylint: disable=protected-access
+        self_attn_forward
+    )
+    self_attn_forward, _ = detab_code(self_attn_forward)
+    assert ORIGINAL_QKV_CODE in self_attn_forward, "Original qkv code not found"
+    assert ORIGINAL_O_CODE in self_attn_forward, "Original o code not found"
+
+    self_attn_forward = self_attn_forward.replace(ORIGINAL_QKV_CODE, PATCHED_QKV_CODE)
+    self_attn_forward = self_attn_forward.replace(ORIGINAL_O_CODE, PATCHED_O_CODE)
+    self_attn_forward = self_attn_forward.replace(
+        "def forward(",
+        "def unsloth_attn_forward(",
+        1,
+    )
+
+    # load imports necessary
+    import transformers.models.llama.modeling_llama
+
+    items_to_import = []
+    for item in dir(transformers.models.llama.modeling_llama):
+        if item in self_attn_forward:
+            items_to_import.append(item)
+
+    exec(  # pylint: disable=exec-used  # nosec B102
+        "from transformers.models.llama.modeling_llama import ("
+        + ", ".join(x for x in items_to_import)
+        + ")",
+        globals(),
+    )
+    exec(self_attn_forward, globals())  # pylint: disable=exec-used  # nosec B102
+    print("patching unsloth attn lora")
+    LlamaFlashAttention2.forward = (
+        unsloth_attn_forward  # pylint: disable=undefined-variable  # noqa: F821
+    )
+
+
+def integrate_lora_mlp_patch(peft_model: PeftModelForCausalLM):
+    if peft_model.base_model.config.model_type in ["llama", "mistral"]:
+        from unsloth.kernels import apply_lora_mlp_swiglu
+
+        apply_lora_mlp = apply_lora_mlp_swiglu
+    elif peft_model.base_model.config.model_type == "gemma":
+        from unsloth.kernels import apply_lora_mlp_geglu_approx
+
+        apply_lora_mlp = apply_lora_mlp_geglu_approx
+    else:
+        raise NotImplementedError(
+            f"Model type {peft_model.base_model.config.model_type} not supported"
+        )
+
+    for idx, layer in enumerate(peft_model.model.model.layers):
+        layer_modules = [
+            getattr(layer.mlp, linear_proj)
+            for linear_proj in ["gate_proj", "up_proj", "down_proj"]
+        ]
+        is_mlp_lora = all(hasattr(module, "lora_A") for module in layer_modules)
+        mlp_no_bias = all(
+            getattr(module, "base_layer", module).bias is None
+            for module in layer_modules
+        )
+        mlp_not_dora = all(
+            getattr(module, "lora_magnitude_vector", None) is None
+            for module in layer_modules
+        )
+
+        if is_mlp_lora and mlp_no_bias and mlp_not_dora:
+            layer.mlp.forward = types.MethodType(apply_lora_mlp, layer.mlp)
+        else:
+            logging.warning("unable to apply unsloth lora mlp patch to layer %d", idx)
+
+
+def integrate_lora_patch(peft_model: PeftModelForCausalLM, cfg):
+    from unsloth.kernels import apply_lora_o, apply_lora_qkv
+
+    for idx, layer in enumerate(peft_model.model.model.layers):
+        if cfg.unsloth_lora_qkv:
+            layer_modules = [
+                getattr(layer.self_attn, linear_proj)
+                for linear_proj in ["q_proj", "k_proj", "v_proj"]
+            ]
+            is_qkv_lora = all(hasattr(module, "lora_A") for module in layer_modules)
+            qkv_no_bias = all(
+                getattr(module, "base_layer", module).bias is None
+                for module in layer_modules
+            )
+            qkv_not_dora = all(
+                getattr(module, "lora_magnitude_vector", None) is None
+                for module in layer_modules
+            )
+
+            if is_qkv_lora and qkv_no_bias and qkv_not_dora:
+                layer.self_attn.apply_qkv = apply_lora_qkv
+            else:
+                layer.self_attn.apply_qkv = original_apply_qkv
+                logging.warning(
+                    "unable to apply unsloth lora qkv patch to layer %d", idx
+                )
+        if cfg.unsloth_lora_o:
+            layer_modules = [
+                getattr(layer.self_attn, linear_proj) for linear_proj in ["o_proj"]
+            ]
+            is_o_lora = all(hasattr(module, "lora_A") for module in layer_modules)
+            o_no_bias = all(
+                getattr(module, "base_layer", module).bias is None
+                for module in layer_modules
+            )
+            o_not_dora = all(
+                getattr(module, "lora_magnitude_vector", None) is None
+                for module in layer_modules
+            )
+
+            if is_o_lora and o_no_bias and o_not_dora:
+                layer.self_attn.apply_o = apply_lora_o
+            else:
+                layer.self_attn.apply_o = original_apply_o
+                logging.warning(
+                    "unable to apply unsloth lora o_proj patch to layer %d", idx
+                )

src/axolotl/prompt_strategies/chat_template.py

@@ -1,24 +1,56 @@
 """
 HF Chat Templates prompt strategy
 """
-from typing import Any, Dict, Optional
+
+import logging
+from typing import Any, Dict, List, Optional
 
 from axolotl.prompt_tokenizers import PromptTokenizingStrategy
 from axolotl.prompters import Prompter
 from axolotl.utils.chat_templates import chat_templates
 
+LOG = logging.getLogger("axolotl")
+
 
 class ChatTemplatePrompter(Prompter):
     """prompter for HF chat templates"""
 
-    def __init__(self, tokenizer, chat_template=None, max_length=2048):
+    def __init__(
+        self,
+        tokenizer,
+        chat_template=None,
+        max_length=2048,
+        message_field_role: str = "from",
+        message_field_content: str = "value",
+        roles: Optional[Dict[str, List[str]]] = None,
+    ):
+        if roles:
+            self.roles = {s: t for t, sources in roles.items() for s in sources}
+        else:
+            self.roles = {
+                "human": "user",
+                "user": "user",
+                "assistant": "assistant",
+                "gpt": "assistant",
+                "system": "system",
+            }
+        self.message_field_role = message_field_role
+        self.message_field_content = message_field_content
         self.tokenizer = tokenizer
         self.chat_template = chat_template
         self.max_length = max_length
 
     def build_prompt(self, conversation, add_generation_prompt=False):
+        turns = [
+            {
+                "role": self.roles[t[self.message_field_role]],
+                "content": t[self.message_field_content],
+            }
+            for t in conversation
+        ]
+
         return self.tokenizer.apply_chat_template(
-            conversation,
+            turns,
             truncation=True,
             max_length=self.max_length,
             add_generation_prompt=add_generation_prompt,
@@ -31,9 +63,19 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
     Tokenizing strategy for instruction-based prompts.
     """
 
+    _messages = "conversations"
+
+    @property
+    def messages(self):
+        return self._messages
+
+    @messages.setter
+    def messages(self, messages):
+        self._messages = messages
+
     def tokenize_prompt(self, prompt):
         turns = self.get_conversation_thread(prompt)
-        prompt_ids = self.prompter.build_prompt([turns[0]], add_generation_prompt=True)
+        prompt_ids = self.prompter.build_prompt(turns[:-1], add_generation_prompt=True)
         input_ids = self.prompter.build_prompt(turns)
 
         if not self.train_on_inputs:
@@ -51,28 +93,37 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
         return tokenized_prompt
 
     def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
-        # remap roles - allow for assistant turn
-        role_map = {
-            "human": "user",
-            "user": "user",
-            "assistant": "assistant",
-            "gpt": "assistant",
-        }
-        turns = [
-            {"role": role_map[t["from"]], "content": t["value"]} for t in conversations
-        ]
-        return turns
+        return prompt[self.messages]
 
 
 def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
     chat_template = (
         ds_cfg["chat_template"] if ds_cfg and "chat_template" in ds_cfg else "chatml"
     )
+    message_field_role = (
+        ds_cfg["message_field_role"]
+        if ds_cfg and "message_field_role" in ds_cfg
+        else "from"
+    )
+    message_field_content = (
+        ds_cfg["message_field_content"]
+        if ds_cfg and "message_field_content" in ds_cfg
+        else "value"
+    )
+    roles = ds_cfg["roles"] if ds_cfg and "roles" in ds_cfg else None
+
     strategy = ChatTemplateStrategy(
-        ChatTemplatePrompter(tokenizer, chat_templates(chat_template)),
+        ChatTemplatePrompter(
+            tokenizer,
+            chat_templates(chat_template),
+            message_field_role=message_field_role,
+            message_field_content=message_field_content,
+            roles=roles,
+        ),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,
     )
+    if ds_cfg and "field_messages" in ds_cfg and hasattr(strategy, "messages"):
+        strategy.messages = ds_cfg["field_messages"]
     return strategy

src/axolotl/prompt_strategies/dpo/llama3.py

@@ -0,0 +1,133 @@
+"""
+DPO strategies for llama-3 chat template
+"""
+
+
+def argilla(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = f"{sample['chosen_response']}<|eot_id|>"
+        sample["rejected"] = f"{sample['rejected_response']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def argilla_chat(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for argilla/dpo-mix-7k conversations
+    """
+
+    def transform_fn(sample):
+        sample[
+            "prompt"
+        ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['chosen'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = f"{sample['chosen'][1]['content']}<|eot_id|>"
+        sample["rejected"] = f"{sample['rejected'][1]['content']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def icr(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    chatml transforms for datasets with system, input, chosen, rejected
+    ex. https://huggingface.co/datasets/argilla/distilabel-intel-orca-dpo-pairs
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['input']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
+        sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    For Intel Orca DPO Pairs
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
+        sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def prompt_pairs(
+    cfg, **kwargs
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = f"{sample['chosen']}<|eot_id|>"
+        sample["rejected"] = f"{sample['rejected']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for ultrafeedback binarized conversations
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["chosen"] = f"{sample['chosen'][1]['content']}<|eot_id|>"
+        sample["rejected"] = f"{sample['rejected'][1]['content']}<|eot_id|>"
+        return sample
+
+    return transform_fn

src/axolotl/prompt_strategies/kto/__init__.py

@@ -0,0 +1,9 @@
+"""
+module for KTO style dataset transform strategies
+"""
+
+from functools import partial
+
+from ..base import load as load_base
+
+load = partial(load_base, module_base="axolotl.prompt_strategies.kto")

src/axolotl/prompt_strategies/kto/chatml.py

@@ -0,0 +1,105 @@
+"""
+KTO strategies for chatml
+"""
+# pylint: disable=duplicate-code
+
+
+def argilla(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['instruction']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def argilla_chat(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for argilla/kto-mix-15k conversations
+    """
+
+    def transform_fn(sample):
+        sample[
+            "prompt"
+        ] = f"<|im_start|>user\n{sample['chosen'][0]['content']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["completion"] = f"{sample['completion'][1]['content']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    For Intel Orca KTO
+    ex: argilla/distilabel-intel-orca-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['question']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def prompt_pairs(
+    cfg, **kwargs
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn
+
+
+def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for ultrafeedback binarized conversations
+    ex: argilla/ultrafeedback-binarized-preferences-cleaned-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|im_start|>system\n{sample['system']}<|im_end|>\n"
+                f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|im_start|>user\n{sample['prompt']}<|im_end|>\n<|im_start|>assistant\n"
+        sample["completion"] = f"{sample['completion']}<|im_end|>"
+        return sample
+
+    return transform_fn

src/axolotl/prompt_strategies/kto/llama3.py

@@ -0,0 +1,105 @@
+"""
+KTO strategies for llama-3 chat template
+"""
+# pylint: disable=duplicate-code
+
+
+def argilla(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['instruction']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def argilla_chat(
+    cfg,
+    **kwargs,
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for argilla/kto-mix-15k conversations
+    """
+
+    def transform_fn(sample):
+        sample[
+            "prompt"
+        ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['completion'][0]['content']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["completion"] = f"{sample['completion'][1]['content']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def intel(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    For Intel Orca KTO
+    ex: argilla/distilabel-intel-orca-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['question']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def prompt_pairs(
+    cfg, **kwargs
+):  # pylint: disable=possibly-unused-variable,unused-argument
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn
+
+
+def ultra(cfg, **kwargs):  # pylint: disable=possibly-unused-variable,unused-argument
+    """
+    for ultrafeedback binarized conversations
+    ex: argilla/ultrafeedback-binarized-preferences-cleaned-kto
+    """
+
+    def transform_fn(sample):
+        if "system" in sample and sample["system"]:
+            sample["prompt"] = (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{sample['system']}<|eot_id|>"
+                f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+            )
+        else:
+            sample[
+                "prompt"
+            ] = f"<|start_header_id|>user<|end_header_id|>\n\n{sample['prompt']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        sample["completion"] = f"{sample['completion']}<|eot_id|>"
+        return sample
+
+    return transform_fn

src/axolotl/prompt_strategies/kto/user_defined.py

@@ -0,0 +1,39 @@
+"""
+User-defined KTO strategies
+"""
+# pylint: disable=duplicate-code
+
+
+def default(cfg, dataset_idx=0, **kwargs):  # pylint: disable=unused-argument
+    ds_cfg = cfg["datasets"][dataset_idx]["type"]
+    if not isinstance(ds_cfg, dict):
+        raise ValueError(
+            f"User-defined dataset type must be a dictionary. Got: {ds_cfg}"
+        )
+    field_prompt = ds_cfg.get("field_prompt", "prompt")
+    field_system = ds_cfg.get("field_system", "system")
+    field_completion = ds_cfg.get("field_completion", "completion")
+    field_label = ds_cfg.get("field_label", "label")
+    prompt_format = ds_cfg.get("prompt_format")
+    if not prompt_format:
+        prompt_format = "{" + field_prompt + "}"
+    completion_format = ds_cfg.get("completion_format")
+    if not completion_format:
+        chosen_format = "{" + field_completion + "}"
+
+    def transform_fn(sample):
+        if (
+            "{" + field_system + "}" in prompt_format
+            and field_system in sample
+            and sample[field_system]
+        ):
+            sample["prompt"] = prompt_format.format(
+                system=sample[field_system], prompt=sample[field_prompt]
+            )
+        else:
+            sample["prompt"] = prompt_format.format(prompt=sample["prompt"])
+        sample["completion"] = chosen_format.format(chosen=sample[field_completion])
+        sample["label"] = sample[field_label]
+        return sample
+
+    return transform_fn

src/axolotl/prompt_strategies/sharegpt.py

@@ -1,7 +1,7 @@
 """Module containing the SimpleShareGPTPromptTokenizingStrategy class"""
 
 import logging
-from typing import Any, Dict, Optional
+from typing import Any, Dict, Optional, Type
 
 from fastchat.conversation import Conversation, SeparatorStyle, register_conv_template
 
@@ -22,7 +22,7 @@ def register_chatml_template(system_message=None):
             name="chatml",
             system_template="<|im_start|>system\n{system_message}",
             system_message=system_message,
-            roles=["<|im_start|>user", "<|im_start|>assistant"],
+            roles=("<|im_start|>user", "<|im_start|>assistant"),
             sep_style=SeparatorStyle.CHATML,
             sep="<|im_end|>",
         )
@@ -32,83 +32,65 @@ def register_chatml_template(system_message=None):
             name="chatml_glaive",
             system_template="<|im_start|>system\n{system_message}",
             system_message=system_message,
-            roles=["<|im_start|>user", "<|im_start|>assistant", "<|im_start|>tool"],
+            roles=("<|im_start|>user", "<|im_start|>assistant", "<|im_start|>tool"),
             sep_style=SeparatorStyle.CHATML,
             sep="<|im_end|>",
         )
     )
 
 
-def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
-    conversation = (
-        ds_cfg["conversation"] if ds_cfg and "conversation" in ds_cfg else None
-    )
-    field_human = ds_cfg["field_human"] if ds_cfg and "field_human" in ds_cfg else None
-    field_model = ds_cfg["field_model"] if ds_cfg and "field_model" in ds_cfg else None
-    roles = ds_cfg["roles"].to_dict() if ds_cfg and "roles" in ds_cfg else None
-    strategy = SimpleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(
-            conversation=conversation,
-            role_key_model=field_model,
-            role_key_human=field_human,
-            roles=roles,
-        ),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-    if ds_cfg and "strict" in ds_cfg:
-        strategy.strict = ds_cfg["strict"]
-    return strategy
-
-
-def load_ultrachat(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
-    conversation = (
-        ds_cfg["conversation"] if ds_cfg and "conversation" in ds_cfg else None
-    )
-    strategy = UltrachatShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(
-            conversation=conversation,
-        ),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
-    if ds_cfg and "strict" in ds_cfg:
-        strategy.strict = ds_cfg["strict"]
-    return strategy
-
-
-def load_role(tokenizer, cfg):
-    return SimpleRoleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
+def register_llama3_template(system_message=None):
+    system_message = system_message or "You are a helpful assistant."
+    register_conv_template(
+        Conversation(
+            name="llama3",
+            system_template="<|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|>",
+            system_message=system_message,
+            roles=("user", "assistant"),
+            sep_style=SeparatorStyle.LLAMA3,
+            sep="",
+            stop_str="<|eot_id|>",
+            stop_token_ids=[128001, 128009],
+        )
     )
 
 
-def load_guanaco(tokenizer, cfg):
-    return GuanacoShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
+def build_loader(
+    tokenization_strategy_cls: Type["ShareGPTPromptTokenizingStrategy"],
+    prompter_cls: Type["ShareGPTPrompterV2"],
+    default_conversation: Optional[str] = None,
+):
+    def _load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
+        conversation = (
+            ds_cfg["conversation"]
+            if ds_cfg and "conversation" in ds_cfg
+            else default_conversation
+        )
+        field_human = (
+            ds_cfg["field_human"] if ds_cfg and "field_human" in ds_cfg else None
+        )
+        field_model = (
+            ds_cfg["field_model"] if ds_cfg and "field_model" in ds_cfg else None
+        )
+        roles = ds_cfg["roles"].to_dict() if ds_cfg and "roles" in ds_cfg else None
+        strategy = tokenization_strategy_cls(
+            prompter_cls(
+                conversation=conversation,
+                role_key_model=field_model,
+                role_key_human=field_human,
+                roles=roles,
+            ),
+            tokenizer,
+            cfg.train_on_inputs,
+            cfg.sequence_len,
+        )
+        if ds_cfg and "strict" in ds_cfg and hasattr(strategy, "strict"):
+            strategy.strict = ds_cfg["strict"]
+        if ds_cfg and "field_messages" in ds_cfg and hasattr(strategy, "messages"):
+            strategy.messages = ds_cfg["field_messages"]
+        return strategy
 
-
-def load_glaive(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
-    conversation = (
-        ds_cfg["conversation"]
-        if ds_cfg and "conversation" in ds_cfg
-        else "chatml_glaive"
-    )
-    return GlaiveShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(conversation=conversation),
-        tokenizer,
-        cfg.train_on_inputs,
-        cfg.sequence_len,
-    )
+    return _load
 
 
 class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
@@ -117,6 +99,7 @@ class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
     """
 
     _strict = False
+    _messages = "conversations"
 
     @property
     def strict(self):
@@ -126,8 +109,16 @@ class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
     def strict(self, strict):
         self._strict = strict
 
+    @property
+    def messages(self):
+        return self._messages
+
+    @messages.setter
+    def messages(self, messages):
+        self._messages = messages
+
     def get_conversation_thread(self, prompt):
-        conversations = prompt["conversations"]
+        conversations = prompt[self.messages]
         if self.strict:
             return conversations
         role_key = "from"
@@ -158,7 +149,9 @@ class SimpleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
         return turns
 
 
-class SimpleRoleShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
+class SimpleRoleShareGPTPromptTokenizingStrategy(
+    SimpleShareGPTPromptTokenizingStrategy
+):
     """
     basic sharegpt strategy to grab conversations from the sample row, but uses role instead of from
     """
@@ -209,3 +202,16 @@ class GlaiveShareGPTPromptTokenizingStrategy(SimpleShareGPTPromptTokenizingStrat
         conversation = merge_consecutive_messages(conversation)
 
         return conversation
+
+
+load = build_loader(SimpleShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2)
+load_role = build_loader(SimpleRoleShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2)
+load_ultrachat = build_loader(
+    UltrachatShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2
+)
+load_guanaco = build_loader(GuanacoShareGPTPromptTokenizingStrategy, ShareGPTPrompterV2)
+load_glaive = build_loader(
+    GlaiveShareGPTPromptTokenizingStrategy,
+    ShareGPTPrompterV2,
+    default_conversation="chatml_glaive",
+)

src/axolotl/prompters.py

@@ -263,6 +263,7 @@ CONVERSATION_ROLE_FORMAT = {
     "chatml": "<|im_start|>{ROLE}",
     "zephyr": "<|{ROLE}|>",
     "vicuna_v1.1": "{ROLE}",
+    "llama3": "<|start_header_id|>{ROLE}<|end_header_id|>",
 }
 
 
@@ -348,7 +349,10 @@ class ShareGPTPrompter(Prompter):  # pylint: disable=too-few-public-methods
                 )
 
             if len(conv.messages) > 0 and ((role == conv.messages[-1][0])):
-                LOG.warning(f"{SHAREGPT_ASSERTION_FAILED_ROLE}: {sentence}")
+                if (
+                    role != "assistant"
+                ):  # back to back assistant calls may be okay for tool calls
+                    LOG.warning(f"{SHAREGPT_ASSERTION_FAILED_ROLE}: {sentence}")
 
             conv.append_message(role, sentence["value"])
 

src/axolotl/train.py

@@ -197,6 +197,13 @@ def train(
         trainer.accelerator.wait_for_everyone()
         unwrapped_model = trainer.accelerator.unwrap_model(trainer.model_wrapped)
 
+        # the trainer saved a model.safetensors file in the output directory,
+        # but it is a proxy model and should be deleted
+        if os.path.exists(os.path.join(cfg.output_dir, "model.safetensors")):
+            LOG.info(f"Deleting {os.path.join(cfg.output_dir, 'model.safetensors')}")
+            LOG.info("This is a proxy model and should be deleted")
+            os.remove(os.path.join(cfg.output_dir, "model.safetensors"))
+
         # Saves the whole/unpartitioned fp16 model when in ZeRO Stage-3 to the output directory if
         # `stage3_gather_16bit_weights_on_model_save` is True in DeepSpeed Config file or
         # `zero3_save_16bit_model` is True in DeepSpeed Plugin.
@@ -212,6 +219,10 @@ def train(
         if cfg.flash_optimum and BetterTransformer:
             model = BetterTransformer.reverse(model)
 
+        if cfg.rl and cfg.adapter and not cfg.rl_adapter_ref_model:
+            trainer.model.save_pretrained(
+                cfg.output_dir, safe_serialization=safe_serialization
+            )
         model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
 
     if not cfg.hub_model_id:

src/axolotl/utils/callbacks/__init__.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import math
 import os
 from shutil import copyfile
 from tempfile import NamedTemporaryFile
@@ -775,9 +776,27 @@ class SaveAxolotlConfigtoWandBCallback(TrainerCallback):
         return control
 
 
-class SaveModelOnTrainEndCallback(TrainerCallback):
+class SaveModelCallback(TrainerCallback):
     """Callback to save model on train end"""
 
+    def on_step_end(  # pylint: disable=unused-argument
+        self,
+        args: TrainingArguments,
+        state: TrainerState,
+        control: TrainerControl,
+        **kwargs,
+    ):
+        # Save
+        if state.global_step >= state.max_steps:
+            control.should_save = True
+        elif (
+            args.save_strategy == IntervalStrategy.STEPS
+            and state.save_steps < 1.0
+            and state.global_step % math.ceil(state.save_steps * state.max_steps) == 0
+        ):
+            # workaround to save model on fractional save_steps
+            control.should_save = True
+
     def on_train_end(  # pylint: disable=unused-argument
         self, args, state, control, **kwargs
     ):

src/axolotl/utils/chat_templates.py

@@ -24,6 +24,7 @@ def chat_templates(user_choice: str):
         "chatml": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
         "gemma": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}",
         "cohere": "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% elif false == true %}{% set loop_messages = messages %}{% set system_message = 'You are Command-R, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses. You are trained by Cohere.' %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% if system_message != false %}{{ '<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>' + system_message + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% set content = message['content'] %}{% if message['role'] == 'user' %}{{ '<|START_OF_TURN_TOKEN|><|USER_TOKEN|>' + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% elif message['role'] == 'assistant' %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>'  + content.strip() + '<|END_OF_TURN_TOKEN|>' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>' }}{% endif %}",
+        "llama3": "{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}",
     }
 
     if user_choice in templates: