chore: update title

.github/workflows/main.yml

@@ -24,7 +24,7 @@ jobs:
             cuda_version: 12.4.1
             python_version: "3.11"
             pytorch: 2.5.1
-            axolotl_extras:
+            axolotl_extras: vllm
           - cuda: 124
             cuda_version: 12.4.1
             python_version: "3.11"

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

@@ -43,7 +43,7 @@ jobs:
             cuda_version: 12.4.1
             python_version: "3.11"
             pytorch: 2.5.1
-            axolotl_extras:
+            axolotl_extras: vllm
             num_gpus: 2
             nightly_build: "true"
           - cuda: 126

.github/workflows/preview-docs.yml

@@ -1,55 +0,0 @@
-name: Preview
-on:
-  workflow_dispatch:
-  pull_request:
-    types: [opened, synchronize, reopened]
-
-permissions:
-  checks: write
-  contents: write
-  deployments: write
-  issues: write
-  discussions: write
-  pages: write
-  pull-requests: write
-  statuses: write
-
-jobs:
-  preview:
-    runs-on: ubuntu-latest
-    steps:
-      - name: Check out repository
-        uses: actions/checkout@v4
-
-      - name: Set up Quarto
-        uses: quarto-dev/quarto-actions/setup@v2
-
-      - name: Setup Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: '3.11'
-
-      - name: Install dependencies
-        run: |
-          python3 -m pip install jupyter quartodoc
-          python3 -m pip install -e . --no-deps
-
-      - name: Build autodoc
-        run: quartodoc build
-
-      - name: Quarto render
-        run: quarto render
-
-      - name: Netlify Publish
-        uses: nwtgck/actions-netlify@v3.0
-        with:
-          publish-dir: './_site'
-          enable-pull-request-comment: true
-          enable-github-deployment: true
-          github-token: ${{ secrets.GITHUB_TOKEN }}
-          deploy-message: "Deployed On Netlify"
-          github-deployment-environment: 'preview'
-          github-deployment-description: 'Preview Deployment'
-        env:
-          NETLIFY_AUTH_TOKEN: ${{ secrets.NETLIFY_AUTH_TOKEN }}
-          NETLIFY_SITE_ID: ${{ secrets.NETLIFY_SITE_ID }}

.github/workflows/tests.yml

@@ -269,7 +269,7 @@ jobs:
             python_version: "3.11"
             pytorch: 2.5.1
             num_gpus: 1
-            axolotl_extras:
+            axolotl_extras: vllm
           - cuda: 126
             cuda_version: 12.6.3
             python_version: "3.11"

.runpod/Dockerfile

@@ -1,10 +1,11 @@
-FROM axolotlai/axolotl-cloud:main-py3.11-cu124-2.6.0
+FROM runpod/pytorch:3.10-2.0.0-117
 
 COPY .runpod/requirements.txt /requirements.txt
 RUN --mount=type=cache,target=/root/.cache/pip \
     python3 -m pip install --upgrade pip && \
     python3 -m pip install --upgrade -r /requirements.txt
 
+
 # Environment settings
 ARG BASE_VOLUME="/runpod-volume"
 ENV BASE_VOLUME=$BASE_VOLUME
@@ -14,5 +15,4 @@ ENV TRANSFORMERS_CACHE="${BASE_VOLUME}/huggingface-cache/hub"
 
 COPY .runpod/src /src
 
-WORKDIR /src
 CMD ["python3", "/src/handler.py"]

.runpod/requirements.txt

@@ -5,3 +5,11 @@
 # git+https://github.com/runpod/runpod-python.git
 # To learn more, see https://pip.pypa.io/en/stable/reference/requirements-file-format/
 runpod~=1.7.0
+huggingface_hub
+typing-extensions
+pydantic
+pydantic-settings
+hf-transfer
+setuptools
+numpy==2.0.0
+axolotl[flash-attn,deepspeed]

.runpod/tests.json

@@ -1,65 +1,69 @@
 {
-  "input": {
-    "name": "quick_smoke_test_sft",
-    "user_id": "user",
-    "model_id": "llama-test",
-    "run_id": "llama-test",
-    "credentials": {
-      "wandb_api_key": "",
-      "hf_token": ""
-    },
-    "args": {
-      "base_model": "HuggingFaceTB/SmolLM2-135M",
-      "model_type": "AutoModelForCausalLM",
-      "tokenizer_type": "AutoTokenizer",
-      "load_in_8bit": true,
-      "load_in_4bit": false,
-      "strict": false,
-      "datasets": [
-        {
-          "path": "mhenrichsen/alpaca_2k_test",
-          "type": "alpaca"
+  "tests": [
+    {
+      "name": "quick_smoke_test_sft",
+      "input": {
+        "user_id": "user",
+        "model_id": "llama-test",
+        "run_id": "llama-test",
+        "credentials": {
+          "wandb_api_key": "",
+          "hf_token": ""
+        },
+        "args": {
+          "base_model": "NousResearch/Meta-Llama-3-8B",
+          "model_type": "LlamaForCausalLM",
+          "tokenizer_type": "AutoTokenizer",
+          "load_in_8bit": true,
+          "load_in_4bit": false,
+          "strict": false,
+          "datasets": [
+            {
+              "path": "mhenrichsen/alpaca_2k_test",
+              "type": "alpaca"
+            }
+          ],
+          "val_set_size": 0.05,
+          "output_dir": "./outputs/lora-out",
+          "sequence_len": 4096,
+          "sample_packing": true,
+          "eval_sample_packing": false,
+          "pad_to_sequence_len": true,
+          "adapter": "lora",
+          "lora_r": 32,
+          "lora_alpha": 16,
+          "lora_dropout": 0.05,
+          "lora_target_linear": true,
+          "lora_modules_to_save": [
+            "embed_tokens",
+            "lm_head"
+          ],
+          "gradient_accumulation_steps": 4,
+          "micro_batch_size": 2,
+          "num_epochs": 1,
+          "optimizer": "adamw_bnb_8bit",
+          "lr_scheduler": "cosine",
+          "learning_rate": 0.0002,
+          "train_on_inputs": false,
+          "group_by_length": false,
+          "bf16": "auto",
+          "tf32": false,
+          "gradient_checkpointing": true,
+          "logging_steps": 1,
+          "flash_attention": true,
+          "warmup_steps": 1,
+          "evals_per_epoch": 1,
+          "eval_max_new_tokens": 128,
+          "saves_per_epoch": 1,
+          "weight_decay": 0.0,
+          "special_tokens": {
+            "pad_token": "<|end_of_text|>"
+          }
         }
-      ],
-      "val_set_size": 0.05,
-      "output_dir": "./outputs/lora-out",
-      "sequence_len": 4096,
-      "sample_packing": true,
-      "eval_sample_packing": false,
-      "pad_to_sequence_len": true,
-      "adapter": "lora",
-      "lora_r": 32,
-      "lora_alpha": 64,
-      "lora_dropout": 0.05,
-      "lora_target_linear": true,
-      "lora_modules_to_save": [
-        "embed_tokens",
-        "lm_head"
-      ],
-      "gradient_accumulation_steps": 4,
-      "micro_batch_size": 2,
-      "num_epochs": 1,
-      "optimizer": "adamw_torch_fused",
-      "lr_scheduler": "cosine",
-      "learning_rate": 0.0002,
-      "train_on_inputs": false,
-      "group_by_length": false,
-      "bf16": "auto",
-      "tf32": true,
-      "gradient_checkpointing": true,
-      "logging_steps": 1,
-      "flash_attention": true,
-      "warmup_steps": 1,
-      "evals_per_epoch": 1,
-      "eval_max_new_tokens": 128,
-      "saves_per_epoch": 1,
-      "weight_decay": 0.0,
-      "special_tokens": {
-        "pad_token": "<|endoftext|>"
-      }
-    },
-    "timeout": 100000
-  },
+      },
+      "timeout": 100000
+    }
+  ],
   "config": {
     "gpuTypeId": "NVIDIA GeForce RTX 4090",
     "gpuCount": 1,

cicd/multigpu.sh

@@ -20,4 +20,4 @@ pytest -v  --durations=10 -n1 /workspace/axolotl/tests/e2e/multigpu/patched/ \
   --cov-report=xml:multigpu-coverage.xml
 
 # Upload coverage to Codecov
-codecov upload-process -t "${CODECOV_TOKEN}" -f multigpu-coverage.xml -F multigpu,docker-tests,pytorch-${PYTORCH_VERSION} || true
+codecov upload-process -t $CODECOV_TOKEN -f multigpu-coverage.xml -F multigpu,docker-tests,pytorch-${PYTORCH_VERSION}

docs/config.qmd

@@ -154,10 +154,6 @@ datasets:
     # Key containing the messages (default: "messages")
     field_messages: messages
 
-    # Key containing the system message (default: "system")
-    # If the system message is not present in the dataset sample, it will be loaded from the field_system property.
-    field_system: system
-
     # Mapping of properties from the input dataset to the chat template.
     # (default: message_property_mappings={'role':'role', 'content':'content'})
     # If a property exists in the template but not in this mapping, the system will attempt
@@ -187,7 +183,7 @@ datasets:
     # IMPORTANT: The following fields determine which parts of the conversation to train on.
     # Priority order: message_field_training > message_field_training_detail > train_on_inputs or role in roles_to_train
     # See examples at `docs/dataset-formats/conversation.qmd`
-    # Note: If the below 5 fields are empty, defaults to training only on the last message.
+    # Note: If the below 4 fields are set to empty, defaults to training only on the last message.
 
     # Optional[List[str]]. Roles to train on. The tokens from these roles will be considered for the loss.
     roles_to_train: ["assistant"]  # default
@@ -196,13 +192,7 @@ datasets:
     # - turn (default): train on the EOS token at the end of each trainable turn
     # - last: train on the last EOS token in the conversation
     # TIP: Please make sure that your `tokenizer.eos_token` is same as EOS/EOT token in template. Otherwise, set `eos_token` under `special_tokens`.
-    train_on_eos: turn
-    # Optional[str]. Which EOT (End-of-Turn) tokens to train on in the conversation. Possible values are:
-    # - all: train on all EOT tokens
-    # - turn: train on the EOT token at the end of each trainable turn
-    # - last: train on the last EOT token in the conversation
-    # If not specified, defaults to the value of train_on_eos for backward compatibility.
-    train_on_eot:
+    train_on_eos: last
     # The key in the message turn that indicates via boolean whether tokens of a turn should be considered for training. Useful to selectively train on certain turns besides the `roles_to_train`.
     message_field_training: training
     # The key in the message turn that contains the training details. Useful to selectively train on certain tokens in a turn.
@@ -285,17 +275,8 @@ process_reward_model:
 chat_template: tokenizer_default
 # custom jinja template for chat template. This will be only used if chat_template is set to `jinja` or `null` (in which case chat_template is automatically set to `jinja`). Default is null.
 chat_template_jinja: null
-# Optional[List[str]]. Custom EOT (End-of-Turn) tokens to mask/unmask during training.
-# These tokens mark the boundaries between conversation turns.
-# For example: ["/INST", "</s>", "[/SYSTEM_PROMPT]"]
-# If not specified, defaults to just the model's eos_token.
-# This is useful for templates that use multiple delimiter tokens.
-eot_tokens:
-  # - "</s>"
-  # - "[/INST]"
-  # - "[/SYSTEM_PROMPT]"
-# Changes the default system message
-default_system_message: You are a helpful assistant. Please give a long and detailed answer. # Currently only supports chatml.
+# Changes the default system message. Currently only supports chatml.
+default_system_message: You are a helpful assistant. Please give a long and detailed answer.
 # Axolotl attempts to save the dataset as an arrow after packing the data together so
 # subsequent training attempts load faster, relative path
 dataset_prepared_path: data/last_run_prepared
@@ -680,10 +661,8 @@ special_tokens:
   # unk_token: "<unk>"
   # pad_token: "[PAD]"
 
-# Optional[list[str]]. Add extra tokens to the tokenizer.
+# Add extra tokens.
 tokens:
-  # - "<|startoftext|>"
-  # - "<|endoftext|>"
 
 # Mapping token_id to new_token_string to override reserved added_tokens in the tokenizer.
 # Only works for tokens that are not part of the base vocab (aka are added_tokens).

docs/dataset-formats/conversation.qmd

@@ -4,6 +4,18 @@ description: Conversation format for supervised fine-tuning.
 order: 3
 ---
 
+## sharegpt
+
+::: {.callout-important}
+ShareGPT is deprecated!. Please see [chat_template](#chat_template) section below.
+:::
+
+## pygmalion
+
+```{.json filename="data.jsonl"}
+{"conversations": [{"role": "...", "value": "..."}]}
+```
+
 ## chat_template
 
 Chat Template strategy uses a jinja2 template that converts a list of messages into a prompt. Support using tokenizer's template, a supported template, or custom jinja2.
@@ -52,7 +64,7 @@ We recommend checking the below examples for other usecases.
 
 ### Examples
 
-1. (Legacy) Using the default chat template in the tokenizer_config.json on OpenAI messages format, training on only last message.
+1. Using the default chat template in the tokenizer_config.json on OpenAI messages format, training on only last message.
 
 ```yaml
 datasets:
@@ -97,55 +109,10 @@ datasets:
 ```
 
 ::: {.callout-important}
-Please make sure that your `tokenizer.eos_token` is same as EOS (End-of-Sequence) token in template. Otherwise, set `eos_token` under `special_tokens: `.
+Please make sure that your `tokenizer.eos_token` is same as EOS/EOT token in template. Otherwise, set `eos_token` under `special_tokens`.
 :::
 
-5. If you are using a template that has a different EOT (End-of-Turn) token from EOS token or multiple EOT tokens (like Mistral V7 Tekken), set the `eot_tokens: ` config. The handling of EOT tokens follows `train_on_eos: ` which defaults to turn.
-
-```yaml
-eot_tokens:
-  - "[/INST]"
-  # - "[/SYSTEM_PROMPT]"
-
-datasets:
-  - path: ...
-    type: chat_template
-
-    # optional
-    train_on_eot: turn  # defaults read from train_on_eos (which defaults to turn)
-```
-
-::: {.callout-tip}
-See [config documentation](../config.qmd) for detailed explanations of "turn", "last", and "all" options for training on tokens.
-:::
-
-::: {.callout-note}
-Using `eot_tokens` requires each token that exists in `chat_template` to be a single token in the tokenizer. Otherwise, the tokenizer will split the token and cause unexpected behavior.
-
-You can add those tokens as new tokens under `tokens: ` or (recommended) override unused added_tokens via `added_tokens_overrides: `. See [config](../config.qmd) for more details.
-:::
-
-6. Continuing from the previous example, if you want to train on all EOT token trainable turns but only last EOS token, set `train_on_eos: last`.
-
-```yaml
-eot_tokens:
-  - "[/INST]"
-  # ...
-
-datasets:
-  - path: ...
-    type: chat_template
-
-    train_on_eos: last
-    train_on_eot: turn
-```
-
-::: {.callout-tip}
-If EOS token only appears at the end of a prompt, `train_on_eos: last` is equivalent to `train_on_eos: turn`. Therefore, generally, you can leave them to their defaults and omit them.
-:::
-
-
-7. (Advanced) Using fine-grained control over tokens and turns to train in a conversation
+5. (Advanced) Using fine-grained control over tokens and turns to train in a conversation
 
 For a data sample that looks like:
 
@@ -195,15 +162,3 @@ datasets:
 ::: {.callout-tip}
 It is not necessary to set both `message_field_training` and `message_field_training_detail` at once.
 :::
-
-## sharegpt
-
-::: {.callout-important}
-ShareGPT is deprecated!. Please see [chat_template](#chat_template) section.
-:::
-
-## pygmalion
-
-```{.json filename="data.jsonl"}
-{"conversations": [{"role": "...", "value": "..."}]}
-```

docs/faq.qmd

@@ -73,40 +73,10 @@ description: Frequently asked questions
 
 > A: This is likely an empty turn.
 
-**Q: The EOS token is incorrectly being masked or not being masked / `EOS token __ not found in chat template`.**
+**Q: The EOS/EOT token is incorrectly being masked or not being masked.**
 
-> A: There can be two reasons:
-
-> 1. This is because of the mismatch between `tokenizer.eos_token` and EOS token in template. Please make sure to set `eos_token: ` under `special_tokens: ` to the same EOS token as in template.
-
-> 2. The EOS token is not in the template. Please check if your template is correct. As an example, `phi_35` template does not use its dedicated EOS token `<|endoftext|>` at the end.
+> A: This is because of the mismatch between `tokenizer.eos_token` and EOS/EOT token in template. Please make sure to set `eos_token` under `special_tokens` to the same EOS/EOT token as in template.
 
 **Q: "`chat_template` choice is `tokenizer_default` but tokenizer's `chat_template` is null. Please add a `chat_template` in tokenizer config"**
 
 > A: This is because the tokenizer does not have a chat template. Please add a chat template in the tokenizer config. See [chat_template](dataset-formats/conversation.qmd#chat-template) for more details.
-
-**Q: The EOT token(s) are incorrectly being masked or not being masked / `EOT token __ not found in chat template`.**
-
-> A: There can be two reasons:
-
-> 1. The EOT token is different from the EOS token and was not specified under `eot_tokens: `. Please set `eot_tokens: ` to the same EOT token(s) as in template.
-
-> 2. There is more than one EOT token per turn in the template. Please raise an issue with examples as we recognize this as an edge case.
-
-**Q: `EOT token encoding failed. Please check if the token is valid and can be encoded.`**
-
-> A: There could be some issue with the tokenizer or unicode encoding. Please raise an issue with examples with the EOT token & tokenizer causing the issue.
-
-**Q: `EOT token __ is encoded as multiple tokens.`**
-
-> A: This is because the EOT token is encoded as multiple tokens which can cause unexpected behavior. Please add it under `tokens: ` or (recommended) override unused added_tokens via `added_tokens_overrides: `.
-
-**Q: `Conflict between train_on_eos and train_on_eot. eos_token is in eot_tokens and train_on_eos != train_on_eot`**
-
-> A: This is because the EOS token is in the `eot_tokens: ` while mismatch between `train_on_eos: ` and `train_on_eot: `. This will cause one to override the other. Please ensure that `train_on_eos: ` and `train_on_eot: ` are the same or remove the EOS token from `eot_tokens: `.
-
-**Q: If `eot_tokens: ` is not provided, what happens?**
-
-> A: If `eot_tokens: ` is not provided, the default behavior is the same as before. EOS tokens used to delimit turns are masked/unmasked depending on whether the turn is trainable.
-
-> Internally, `eot_tokens: tokenizer.eos_token` and `train_on_eot: train_on_eos` (which defaults to `turn`). This transition helps clarify the naming and behavior of EOT/EOS tokens.

docs/rlhf.qmd

@@ -502,7 +502,9 @@ The input format is a simple JSON input with customizable fields based on the ab
 Check out our [GRPO cookbook](https://github.com/axolotl-ai-cloud/axolotl-cookbook/tree/main/grpo#training-an-r1-style-large-language-model-using-grpo).
 :::
 
-In the latest GRPO implementation, `vLLM` is used to significantly speedup trajectory generation during training. In this example, we're using 4 GPUs - 2 for training, and 2 for vLLM:
+If you have multiple GPUs available, we reccomend using `vLLM` with the `GRPOTrainer` to significantly speedup trajectory generation during training.
+First, launch a `vLLM` server using `trl vllm-serve` - you may use a config file or CLI overrides to configure your vLLM server. In this example, we're
+using 4 GPUs - 2 for training, and 2 for vLLM:
 
 ::: {.callout-important}
 Make sure you've installed the correct version of vLLM by including it as an extra when installing axolotl, e.g. `pip install axolotl[vllm]`.
@@ -537,10 +539,6 @@ Your `vLLM` instance will now attempt to spin up, and it's time to kick off trai
 CUDA_VISIBLE_DEVICES=0,1 axolotl train grpo.yaml --num-processes 2
 ```
 
-::: {.callout-note}
-Due to TRL's implementation with vLLM, the vLLM instance must use the last N GPUs instead of the first N GPUs. This is why in the example above, we use `CUDA_VISIBLE_DEVICES=2,3` for the vLLM instance.
-:::
-
 #### Reward functions
 
 GRPO uses custom reward functions and transformations. Please have them ready locally.

examples/qwen3/32b-qlora.yaml

@@ -1,69 +0,0 @@
-base_model: Qwen/Qwen3-32B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-plugins:
-  - axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
-strict: false
-
-chat_template: qwen3
-datasets:
-  - path: mlabonne/FineTome-100k
-    type: chat_template
-    split: train[:20%]
-    field_messages: conversations
-    message_property_mappings:
-      role: from
-      content: value
-val_set_size: 0.0
-output_dir: ./outputs/out
-dataset_prepared_path: last_run_prepared
-
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: true
-pad_to_sequence_len: true
-
-load_in_4bit: true
-adapter: qlora
-lora_r: 16
-lora_alpha: 32
-lora_target_modules:
-  - q_proj
-  - k_proj
-  - v_proj
-  - o_proj
-  - down_proj
-  - up_proj
-lora_mlp_kernel: true
-lora_qkv_kernel: true
-lora_o_kernel: true
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 2
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_4bit
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: offload
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_steps: 10
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-special_tokens:

examples/qwen3/qlora-fsdp.yaml

@@ -1,68 +0,0 @@
-base_model: Qwen/Qwen3-8B
-# Automatically upload checkpoint and final model to HF
-# hub_model_id: username/custom_model_name
-
-load_in_8bit: false
-load_in_4bit: true
-strict: false
-
-datasets:
-  - path: tatsu-lab/alpaca
-    type: alpaca
-dataset_prepared_path:
-val_set_size: 0.05
-output_dir: ./outputs/out
-
-sequence_len: 2048
-sample_packing: true
-eval_sample_packing: true
-pad_to_sequence_len: true
-
-adapter: qlora
-lora_model_dir:
-lora_r: 32
-lora_alpha: 64
-lora_dropout: 0.05
-lora_target_linear: true
-
-wandb_project:
-wandb_entity:
-wandb_watch:
-wandb_name:
-wandb_log_model:
-
-gradient_accumulation_steps: 4
-micro_batch_size: 1
-num_epochs: 1
-optimizer: adamw_torch_fused
-lr_scheduler: cosine
-learning_rate: 0.0002
-
-bf16: auto
-tf32: true
-
-gradient_checkpointing: true
-gradient_checkpointing_kwargs:
-  use_reentrant: false
-resume_from_checkpoint:
-logging_steps: 1
-flash_attention: true
-
-warmup_steps: 10
-evals_per_epoch: 4
-saves_per_epoch: 1
-weight_decay: 0.0
-fsdp:
-  - full_shard
-  - auto_wrap
-fsdp_config:
-  fsdp_limit_all_gathers: true
-  fsdp_sync_module_states: true
-  fsdp_offload_params: true
-  fsdp_use_orig_params: false
-  fsdp_cpu_ram_efficient_loading: true
-  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
-  fsdp_transformer_layer_cls_to_wrap: Qwen3DecoderLayer
-  fsdp_state_dict_type: FULL_STATE_DICT
-  fsdp_sharding_strategy: FULL_SHARD
-special_tokens:

requirements.txt

@@ -11,13 +11,13 @@ liger-kernel==0.5.8
 
 packaging==23.2
 
-peft==0.15.2
+peft==0.15.1
 transformers==4.51.3
 tokenizers>=0.21.1
 accelerate==1.6.0
 datasets==3.5.0
 deepspeed>=0.15.4
-trl==0.17.0
+trl==0.16.1
 hf_xet==1.0.0
 hqq==0.2.5
 

setup.py

@@ -67,13 +67,13 @@ def parse_requirements(extras_require_map):
             if (major, minor) >= (2, 7):
                 _install_requires.pop(_install_requires.index(xformers_version))
                 # _install_requires.append("xformers==0.0.29.post3")  # xformers seems to be hard pinned to 2.6.0
-                extras_require_map["vllm"] = ["vllm==0.8.4"]
+                extras_require_map["vllm"] = ["vllm==0.8.3"]
             elif (major, minor) >= (2, 6):
                 _install_requires.pop(_install_requires.index(xformers_version))
                 _install_requires.append(
                     "xformers==0.0.29.post2"
                 )  # vllm needs post2 w torch 2.6
-                extras_require_map["vllm"] = ["vllm==0.8.4"]
+                extras_require_map["vllm"] = ["vllm==0.8.3"]
             elif (major, minor) >= (2, 5):
                 _install_requires.pop(_install_requires.index(xformers_version))
                 if patch == 0:

src/axolotl/__init__.py

@@ -4,4 +4,4 @@ import pkgutil
 
 __path__ = pkgutil.extend_path(__path__, __name__)  # Make this a namespace package
 
-__version__ = "0.9.0"
+__version__ = "0.8.0"

src/axolotl/cli/train.py

@@ -1,6 +1,5 @@
 """CLI to run training on a model."""
 
-import gc
 import logging
 import os
 from pathlib import Path
@@ -49,11 +48,8 @@ def do_train(cfg: DictDefault, cli_args: TrainerCliArgs):
         dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
 
     model, tokenizer, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
-
     del model, tokenizer, trainer
 
-    gc.collect()
-
     plugin_manager = PluginManager.get_instance()
     plugin_manager.post_train_unload(cfg)
 

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

@@ -3,29 +3,15 @@ DPO trainer for axolotl
 """
 
 import gc
-import random
 from functools import wraps
-from typing import Any, Dict, Optional, Union
+from typing import Any, Dict, Union
 
-import pandas as pd
 import torch
-import wandb
-from accelerate import PartialState
-from datasets import Dataset, IterableDataset
 from peft.optimizers import create_loraplus_optimizer
 from torch import nn
-from torch.utils.data import DataLoader
-from transformers import (
-    BaseImageProcessor,
-    FeatureExtractionMixin,
-    PreTrainedTokenizerBase,
-    ProcessorMixin,
-    Trainer,
-)
-from transformers.trainer_utils import EvalLoopOutput
+from transformers import Trainer
 from transformers.utils import is_sagemaker_mp_enabled
-from trl import DPOConfig, DPOTrainer, maybe_apply_chat_template, maybe_extract_prompt
-from trl.trainer.utils import log_table_to_comet_experiment
+from trl import DPOTrainer
 
 from axolotl.core.trainers.mixins import RngLoaderMixin, SchedulerMixin
 from axolotl.core.trainers.utils import (
@@ -95,64 +81,6 @@ class AxolotlDPOTrainer(RngLoaderMixin, SchedulerMixin, DPOTrainer):
 
         return super().push_to_hub(*args, **kwargs)
 
-    # TODO: remove this once https://github.com/huggingface/trl/pull/3377 is in a release
-    def _prepare_dataset(
-        self,
-        dataset: Union[Dataset, IterableDataset],
-        processing_class: Union[
-            PreTrainedTokenizerBase,
-            BaseImageProcessor,
-            FeatureExtractionMixin,
-            ProcessorMixin,
-        ],
-        args: DPOConfig,
-        dataset_name: str,
-    ) -> Union[Dataset, IterableDataset]:
-        # Build the kwargs for the `map` function
-        map_kwargs: Dict[str, Any] = {"writer_batch_size": 10}
-        if isinstance(dataset, Dataset):  # IterableDataset does not support num_proc
-            map_kwargs["num_proc"] = args.dataset_num_proc
-
-        with PartialState().main_process_first():
-            # Extract prompt if needed
-            if isinstance(
-                dataset, Dataset
-            ):  # `IterableDataset.map` does not support `desc`
-                map_kwargs["desc"] = f"Extracting prompt in {dataset_name} dataset"
-            dataset = dataset.map(maybe_extract_prompt, **map_kwargs)
-
-            # Apply the chat template if needed
-            if isinstance(
-                dataset, Dataset
-            ):  # `IterableDataset.map` does not support `desc`
-                map_kwargs["desc"] = f"Applying chat template to {dataset_name} dataset"
-            dataset = dataset.map(
-                maybe_apply_chat_template,
-                fn_kwargs={"tokenizer": processing_class, "tools": args.tools},
-                **map_kwargs,
-            )
-
-            # Tokenize the dataset
-            if isinstance(
-                dataset, Dataset
-            ):  # `IterableDataset.map` does not support `desc`
-                map_kwargs["desc"] = f"Tokenizing {dataset_name} dataset"
-
-            dataset = dataset.map(
-                self.tokenize_row if not self.is_vision_model else self.process_row,
-                remove_columns=["chosen", "rejected"],
-                fn_kwargs={
-                    "processing_class": processing_class,
-                    "max_prompt_length": args.max_prompt_length,
-                    "max_completion_length": args.max_completion_length,
-                    # for enc-dec, we add the special tokens ([bos_token] + prompt + [eos_token]; completion + [eos_token])
-                    "add_special_tokens": False,
-                },
-                **map_kwargs,
-            )
-
-        return dataset
-
     @staticmethod
     def tokenize_row(
         features,
@@ -196,67 +124,3 @@ class AxolotlDPOTrainer(RngLoaderMixin, SchedulerMixin, DPOTrainer):
         gc.collect()
         torch.cuda.empty_cache()
         return loss
-
-    # TODO: remove this once https://github.com/huggingface/trl/pull/3377 is in a release
-    def evaluation_loop(
-        self,
-        dataloader: DataLoader,
-        description: str,
-        prediction_loss_only: Optional[bool] = None,
-        ignore_keys: Optional[list[str]] = None,
-        metric_key_prefix: str = "eval",
-    ) -> EvalLoopOutput:
-        """
-        Overriding built-in evaluation loop to store metrics for each batch.
-        Prediction/evaluation loop, shared by `Trainer.evaluate()` and `Trainer.predict()`.
-
-        Works both with or without labels.
-        """
-
-        # Sample and save to game log if requested (for one batch to save time)
-        if self.generate_during_eval:
-            # Generate random indices within the range of the total number of samples
-            num_samples = len(dataloader.dataset)
-            random_indices = random.sample(
-                range(num_samples), k=self.args.eval_batch_size
-            )
-
-            # Use dataloader.dataset.select to get the random batch without iterating over the DataLoader
-            random_batch_dataset = dataloader.dataset.select(random_indices)
-            random_batch = self.data_collator(random_batch_dataset)
-            random_batch = self._prepare_inputs(random_batch)
-
-            policy_output_decoded, ref_output_decoded = (
-                self.generate_from_model_and_ref(self.model, random_batch)
-            )
-
-            table = pd.DataFrame(
-                columns=["Prompt", "Policy", "Ref Model"],
-                data=[
-                    [prompt, pol[len(prompt) :], ref[len(prompt) :]]
-                    for prompt, pol, ref in zip(
-                        random_batch_dataset["prompt"],
-                        policy_output_decoded,
-                        ref_output_decoded,
-                    )
-                ],
-            )
-            if "wandb" in self.args.report_to and self.accelerator.is_main_process:
-                wandb.log({"game_log": wandb.Table(data=table)})
-
-            if "comet_ml" in self.args.report_to:
-                log_table_to_comet_experiment(
-                    name="game_log.csv",
-                    table=table,
-                )
-
-        # Base evaluation
-        initial_output = super().evaluation_loop(
-            dataloader,
-            description,
-            prediction_loss_only,
-            ignore_keys,
-            metric_key_prefix,
-        )
-
-        return initial_output

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

@@ -135,9 +135,7 @@ class GRPOStrategy:
         try:
             # use importlib to dynamically load the reward function from the module
             reward_func_module_name = reward_func_fqn.split(".")[-1]
-            reward_func_module = importlib.import_module(
-                ".".join(reward_func_fqn.split(".")[:-1])
-            )
+            reward_func_module = importlib.import_module(reward_func_fqn.split(".")[-2])
             reward_func = getattr(reward_func_module, reward_func_module_name)
             if not len(inspect.signature(reward_func).parameters) >= 2:
                 raise ValueError(

src/axolotl/integrations/base.py

@@ -36,10 +36,9 @@ class BasePlugin:
     Methods:
     register(cfg): Registers the plugin with the given configuration.
     pre_model_load(cfg): Performs actions before the model is loaded.
-    post_model_build(cfg, model): Performs actions after the model is loaded, but before LoRA adapters are applied.
+    post_model_load(cfg, model): Performs actions after the model is loaded.
     pre_lora_load(cfg, model): Performs actions before LoRA weights are loaded.
     post_lora_load(cfg, model): Performs actions after LoRA weights are loaded.
-    post_model_load(cfg, model): Performs actions after the model is loaded, inclusive of any adapters.
     create_optimizer(cfg, trainer): Creates and returns an optimizer for training.
     create_lr_scheduler(cfg, trainer, optimizer): Creates and returns a learning rate scheduler.
     add_callbacks_pre_trainer(cfg, model): Adds callbacks to the trainer before training.
@@ -78,14 +77,6 @@ class BasePlugin:
         None
         """
 
-    def post_model_build(self, cfg, model):  # pylint: disable=unused-argument
-        """
-        Performs actions after the model is built/loaded, but before any adapters are applied.
-
-        Args:
-            cfg (dict): The configuration for the plugin.
-        """
-
     def post_model_load(self, cfg, model):  # pylint: disable=unused-argument
         """
         Performs actions after the model is loaded.
@@ -338,22 +329,9 @@ class PluginManager:
         for plugin in self.plugins.values():
             plugin.pre_model_load(cfg)
 
-    def post_model_build(self, cfg, model):
-        """
-        Calls the post_model_build method of all registered plugins after the model has been built/loaded,
-        but before any adapters have been applied.
-
-        Args:
-            cfg (dict): The configuration for the plugins.
-            model (object): The loaded model.
-        """
-        for plugin in self.plugins.values():
-            plugin.post_model_build(cfg, model)
-
     def post_model_load(self, cfg, model):
         """
-        Calls the post_model_load method of all registered plugins after the model has been loaded
-        inclusive of any adapters
+        Calls the post_model_load method of all registered plugins.
 
         Parameters:
         cfg (dict): The configuration for the plugins.
@@ -480,20 +458,6 @@ class PluginManager:
                 callbacks.extend(plugin_callbacks)
         return callbacks
 
-    def post_train(self, cfg, model):
-        """
-        Calls the post_train method of all registered plugins.
-
-        Parameters:
-        cfg (dict): The configuration for the plugins.
-        model (object): The loaded model.
-
-        Returns:
-        None
-        """
-        for plugin in self.plugins.values():
-            plugin.post_train(cfg, model)
-
     def post_train_unload(self, cfg):
         """
         Calls the post_train_unload method of all registered plugins.

src/axolotl/integrations/cut_cross_entropy/README.md

@@ -32,8 +32,8 @@ plugins:
 ## Supported Models
 
 - llama
-- llama4
 - llama4_text
+- llama4
 - mllama
 - phi3
 - gemma
@@ -43,11 +43,6 @@ plugins:
 - mistral
 - mistral3
 - qwen2
-- qwen2_moe
-- qwen2_vl
-- qwen2_5_vl
-- qwen3
-- qwen3_moe
 - cohere
 - cohere2
 - glm

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

@@ -1,174 +0,0 @@
-"""Llama CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-
-from types import MethodType
-from typing import Optional, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-    apply_lce,
-)
-from transformers.cache_utils import Cache
-from transformers.modeling_outputs import (
-    BaseModelOutputWithPast,
-    CausalLMOutputWithPast,
-)
-from transformers.models.llama.modeling_llama import (
-    _CONFIG_FOR_DOC,
-    LLAMA_INPUTS_DOCSTRING,
-    KwargsForCausalLM,
-)
-from transformers.processing_utils import Unpack
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@can_return_tuple
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
-def cce_forward(
-    self,
-    input_ids: Optional[torch.LongTensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[Cache] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    labels: Optional[torch.LongTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    cache_position: Optional[torch.LongTensor] = None,
-    logits_to_keep: Union[int, torch.Tensor] = 0,
-    **kwargs: Unpack[KwargsForCausalLM],
-) -> CausalLMOutputWithPast:
-    r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
-        logits_to_keep (`int` or `torch.Tensor`, *optional*):
-            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
-            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
-            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
-            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
-            This is useful when using packed tensor format (single dimension for batch and sequence length).
-
-    Returns:
-
-    Example:
-
-    ```python
-    >>> from transformers import AutoTokenizer, LlamaForCausalLM
-
-    >>> model = LlamaForCausalLM.from_pretrained("meta-llama/Llama-2-7b-hf")
-    >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf")
-
-    >>> prompt = "Hey, are you conscious? Can you talk to me?"
-    >>> inputs = tokenizer(prompt, return_tensors="pt")
-
-    >>> # Generate
-    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
-    ```"""
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-
-    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-    outputs: BaseModelOutputWithPast = self.model(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        position_ids=position_ids,
-        past_key_values=past_key_values,
-        inputs_embeds=inputs_embeds,
-        use_cache=use_cache,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        cache_position=cache_position,
-        **kwargs,
-    )
-
-    hidden_states = outputs.last_hidden_state
-    if hidden_states is None:
-        raise ValueError("hidden_states is None")
-
-    loss = None
-    logits = None
-
-    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-    slice_indices = (
-        slice(-logits_to_keep, None)
-        if isinstance(logits_to_keep, int)
-        else logits_to_keep
-    )
-    if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
-        assert labels is not None
-        loss = apply_lce(
-            hidden_states[:, slice_indices, :],
-            self.lm_head.weight,
-            labels,
-            _PATCH_OPTS,
-            **kwargs,
-        )
-    else:
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
-
-        if labels is not None:
-            loss = self.loss_function(
-                logits=logits,
-                labels=labels,
-                vocab_size=self.config.vocab_size,
-                **kwargs,
-            )
-
-    return CausalLMOutputWithPast(
-        loss=loss,
-        logits=logits,
-        past_key_values=outputs.past_key_values,
-        hidden_states=outputs.hidden_states,
-        attentions=outputs.attentions,
-    )
-
-
-def patch_llama(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    """Patch Llama for CCE."""
-    global _PATCH_OPTS  # pylint: disable=global-statement
-    from transformers.models.llama import modeling_llama
-
-    _PATCH_OPTS = patch_options
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_llama.LlamaForCausalLM
-        ), f"Expected a LlamaForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward, maybe_model)
-        return maybe_model
-
-    modeling_llama.LlamaForCausalLM.forward = cce_forward
-    return None

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

@@ -5,7 +5,9 @@
 import transformers
 from cut_cross_entropy.cce_utils import LinearCrossEntropyImpl
 from cut_cross_entropy.linear_cross_entropy import LCE_IMPL_DEFAULT
+from cut_cross_entropy.transformers.llama import patch_llama
 from cut_cross_entropy.transformers.phi3 import patch_phi3
+from cut_cross_entropy.transformers.qwen2 import patch_qwen2
 from cut_cross_entropy.transformers.utils import PatchOptions, TransformersModelT
 
 from axolotl.integrations.cut_cross_entropy.monkeypatch.cohere import (
@@ -22,9 +24,6 @@ from axolotl.integrations.cut_cross_entropy.monkeypatch.glm4 import (
     patch_glm,
     patch_glm4,
 )
-from axolotl.integrations.cut_cross_entropy.monkeypatch.llama import (
-    patch_llama,
-)
 from axolotl.integrations.cut_cross_entropy.monkeypatch.llama4 import (
     patch_llama4,
     patch_llama4_text,
@@ -34,22 +33,6 @@ from axolotl.integrations.cut_cross_entropy.monkeypatch.mistral3 import (
     patch_mistral3,
 )
 from axolotl.integrations.cut_cross_entropy.monkeypatch.mllama import patch_mllama
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2 import (
-    patch_qwen2,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2_5_vl import (
-    patch_qwen2_5_vl,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2_moe import (
-    patch_qwen2_moe,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen2_vl import (
-    patch_qwen2_vl,
-)
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen3 import patch_qwen3
-from axolotl.integrations.cut_cross_entropy.monkeypatch.qwen3_moe import (
-    patch_qwen3_moe,
-)
 
 CUT_CROSS_ENTROPY_MODEL_MAPPING = {
     "llama": patch_llama,
@@ -64,11 +47,6 @@ CUT_CROSS_ENTROPY_MODEL_MAPPING = {
     "mistral": patch_mistral,
     "mistral3": patch_mistral3,
     "qwen2": patch_qwen2,
-    "qwen2_moe": patch_qwen2_moe,
-    "qwen2_vl": patch_qwen2_vl,
-    "qwen2_5_vl": patch_qwen2_5_vl,
-    "qwen3": patch_qwen3,
-    "qwen3_moe": patch_qwen3_moe,
     "cohere": patch_cohere,
     "cohere2": patch_cohere2,
     "glm": patch_glm,

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

@@ -1,37 +0,0 @@
-"""Qwen2 CCE patch. The model inherits Llama's modeling code and uses the same forward method."""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-)
-
-
-def patch_qwen2(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    from transformers.models.qwen2 import modeling_qwen2
-
-    # Set the _PATCH_OPTS in the llama patch file
-    import axolotl.integrations.cut_cross_entropy.monkeypatch.llama as llama_patch
-
-    llama_patch._PATCH_OPTS = patch_options  # pylint: disable=protected-access
-
-    from axolotl.integrations.cut_cross_entropy.monkeypatch.llama import (
-        cce_forward,
-    )
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_qwen2.Qwen2ForCausalLM
-        ), f"Expected a Qwen2ForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward, maybe_model)
-        return maybe_model
-
-    modeling_qwen2.Qwen2ForCausalLM.forward = cce_forward
-    return None

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

@@ -1,246 +0,0 @@
-"""Qwen2.5 VL CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-    apply_lce,
-)
-from torch.nn import CrossEntropyLoss
-from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (
-    Qwen2_5_VLCausalLMOutputWithPast,
-)
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-def cce_forward_multimodal(
-    self,
-    input_ids: Optional[torch.LongTensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[list[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    labels: Optional[torch.LongTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    return_dict: Optional[bool] = None,
-    pixel_values: Optional[torch.Tensor] = None,
-    pixel_values_videos: Optional[torch.FloatTensor] = None,
-    image_grid_thw: Optional[torch.LongTensor] = None,
-    video_grid_thw: Optional[torch.LongTensor] = None,
-    rope_deltas: Optional[torch.LongTensor] = None,
-    cache_position: Optional[torch.LongTensor] = None,
-    second_per_grid_ts: Optional[torch.Tensor] = None,
-) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
-    r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
-    Returns:
-
-    Example:
-
-    ```python
-    >>> from PIL import Image
-    >>> import requests
-    >>> from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration
-
-    >>> model = Qwen2_5_VLForConditionalGeneration.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
-    >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
-
-    >>> messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": "What is shown in this image?"},
-            ],
-        },
-    ]
-    >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
-    >>> image = Image.open(requests.get(url, stream=True).raw)
-
-    >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
-
-    >>> # Generate
-    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
-    ```"""
-
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-    return_dict = (
-        return_dict if return_dict is not None else self.config.use_return_dict
-    )
-
-    if inputs_embeds is None:
-        inputs_embeds = self.model.embed_tokens(input_ids)
-        if pixel_values is not None:
-            pixel_values = pixel_values.type(self.visual.dtype)
-            image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
-            n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
-            n_image_features = image_embeds.shape[0]
-            if n_image_tokens != n_image_features:
-                raise ValueError(
-                    f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
-                )
-
-            mask = input_ids == self.config.image_token_id
-            mask_unsqueezed = mask.unsqueeze(-1)
-            mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
-            image_mask = mask_expanded.to(inputs_embeds.device)
-
-            image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)  # type: ignore
-
-        if pixel_values_videos is not None:
-            pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
-            video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
-            n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
-            n_video_features = video_embeds.shape[0]
-            if n_video_tokens != n_video_features:
-                raise ValueError(
-                    f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
-                )
-
-            mask = input_ids == self.config.video_token_id
-            mask_unsqueezed = mask.unsqueeze(-1)
-            mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
-            video_mask = mask_expanded.to(inputs_embeds.device)
-
-            video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)  # type: ignore
-
-        if attention_mask is not None:
-            attention_mask = attention_mask.to(inputs_embeds.device)
-
-    # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
-    if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
-        # calculate RoPE index once per generation in the pre-fill stage only
-        if (
-            (cache_position is not None and cache_position[0] == 0)
-            or self.rope_deltas is None
-            or (past_key_values is None or past_key_values.get_seq_length() == 0)  # type: ignore
-        ):
-            position_ids, rope_deltas = self.get_rope_index(
-                input_ids,
-                image_grid_thw,
-                video_grid_thw,
-                second_per_grid_ts,
-                attention_mask,
-            )
-            self.rope_deltas = rope_deltas
-        # then use the prev pre-calculated rope-deltas to get the correct position ids
-        else:
-            batch_size, seq_length, _ = inputs_embeds.shape
-            delta = (
-                (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
-                if cache_position is not None
-                else 0
-            )
-            position_ids = torch.arange(seq_length, device=inputs_embeds.device)  # type: ignore
-            position_ids = position_ids.view(1, -1).expand(batch_size, -1)  # type: ignore
-            if cache_position is not None:  # otherwise `deltas` is an int `0`
-                delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)  # type: ignore
-            position_ids = position_ids.add(delta)  # type: ignore
-            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)  # type: ignore
-
-    outputs = self.model(
-        input_ids=None,
-        position_ids=position_ids,
-        attention_mask=attention_mask,
-        past_key_values=past_key_values,
-        inputs_embeds=inputs_embeds,
-        use_cache=use_cache,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        return_dict=return_dict,
-        cache_position=cache_position,
-    )
-
-    hidden_states = outputs[0]
-    logits = None
-    loss = None
-
-    if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
-        assert labels is not None
-        loss = apply_lce(
-            hidden_states,
-            self.lm_head.weight,
-            labels,
-            _PATCH_OPTS,
-        )
-    else:
-        logits = self.lm_head(hidden_states)
-
-        if labels is not None:
-            # Upcast to float if we need to compute the loss to avoid potential precision issues
-            logits = logits.float()
-            # 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)
-
-    if not return_dict:
-        output = (logits,) + outputs[1:]
-        return (loss,) + output if loss is not None else output
-
-    return Qwen2_5_VLCausalLMOutputWithPast(
-        loss=loss,
-        logits=logits,
-        past_key_values=outputs.past_key_values,
-        hidden_states=outputs.hidden_states,
-        attentions=outputs.attentions,
-        rope_deltas=self.rope_deltas,
-    )
-
-
-def patch_qwen2_5_vl(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    global _PATCH_OPTS  # pylint: disable=global-statement
-
-    from transformers.models.qwen2_5_vl import modeling_qwen2_5_vl
-
-    _PATCH_OPTS = patch_options
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_qwen2_5_vl.Qwen2_5_VLForConditionalGeneration
-        ), f"Expected a Qwen2_5_VLForConditionalGeneration model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
-        return maybe_model
-
-    modeling_qwen2_5_vl.Qwen2_5_VLForConditionalGeneration.forward = (
-        cce_forward_multimodal
-    )
-    return None

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

@@ -1,188 +0,0 @@
-"""Qwen2 MoE CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-    apply_lce,
-)
-from transformers.models.qwen2_moe.modeling_qwen2_moe import (
-    _CONFIG_FOR_DOC,
-    QWEN2MOE_INPUTS_DOCSTRING,
-    MoeCausalLMOutputWithPast,
-    MoeModelOutputWithPast,
-    load_balancing_loss_func,
-)
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@can_return_tuple
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(QWEN2MOE_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=MoeCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
-def forward(
-    self,
-    input_ids: Optional[torch.LongTensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[list[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    labels: Optional[torch.LongTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    output_router_logits: Optional[bool] = None,
-    cache_position: Optional[torch.LongTensor] = None,
-    logits_to_keep: Union[int, torch.Tensor] = 0,
-    **loss_kwargs,
-) -> MoeCausalLMOutputWithPast:
-    r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
-        logits_to_keep (`int` or `torch.Tensor`, *optional*):
-            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
-            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
-            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
-            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
-            This is useful when using packed tensor format (single dimension for batch and sequence length).
-
-    Returns:
-
-    Example:
-
-    ```python
-    >>> from transformers import AutoTokenizer, Qwen2MoeForCausalLM
-
-    >>> model = Qwen2MoeForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
-    >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
-
-    >>> prompt = "Hey, are you conscious? Can you talk to me?"
-    >>> inputs = tokenizer(prompt, return_tensors="pt")
-
-    >>> # Generate
-    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
-    ```"""
-
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_router_logits = (
-        output_router_logits
-        if output_router_logits is not None
-        else self.config.output_router_logits
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-
-    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-    outputs: MoeModelOutputWithPast = self.model(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        position_ids=position_ids,
-        past_key_values=past_key_values,
-        inputs_embeds=inputs_embeds,
-        use_cache=use_cache,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        output_router_logits=output_router_logits,
-        cache_position=cache_position,
-    )
-
-    hidden_states = outputs.last_hidden_state
-    loss = None
-    logits = None
-
-    if hidden_states is None:
-        raise ValueError("hidden_states is None")
-
-    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-    slice_indices = (
-        slice(-logits_to_keep, None)
-        if isinstance(logits_to_keep, int)
-        else logits_to_keep
-    )
-
-    if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
-        assert labels is not None
-        loss = apply_lce(
-            hidden_states[:, slice_indices, :],
-            self.lm_head.weight,
-            labels,
-            _PATCH_OPTS,
-            **loss_kwargs,
-        )
-    else:
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
-
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
-
-    aux_loss = None
-    if output_router_logits:
-        aux_loss = load_balancing_loss_func(
-            outputs.router_logits,
-            self.num_experts,
-            self.num_experts_per_tok,
-            attention_mask,
-        )
-        if labels is not None:
-            loss += self.router_aux_loss_coef * aux_loss.to(  # type: ignore
-                loss.device  # type: ignore
-            )  # make sure to reside in the same device
-
-    return MoeCausalLMOutputWithPast(
-        loss=loss,
-        aux_loss=aux_loss,  # type: ignore
-        logits=logits,
-        past_key_values=outputs.past_key_values,
-        hidden_states=outputs.hidden_states,
-        attentions=outputs.attentions,
-        router_logits=outputs.router_logits,
-    )
-
-
-def patch_qwen2_moe(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    global _PATCH_OPTS  # pylint: disable=global-statement
-
-    from transformers.models.qwen2_moe import modeling_qwen2_moe
-
-    _PATCH_OPTS = patch_options
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_qwen2_moe.Qwen2MoeForCausalLM
-        ), f"Expected a Qwen3MoeForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(forward, maybe_model)
-
-        return maybe_model
-
-    modeling_qwen2_moe.Qwen2MoeForCausalLM.forward = forward
-    return None

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

@@ -1,249 +0,0 @@
-"""Qwen2 VL CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Tuple, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-    apply_lce,
-)
-from torch.nn import CrossEntropyLoss
-from transformers.models.qwen2_vl.modeling_qwen2_vl import (
-    _CONFIG_FOR_DOC,
-    QWEN2_VL_INPUTS_DOCSTRING,
-    Qwen2VLCausalLMOutputWithPast,
-)
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@add_start_docstrings_to_model_forward(QWEN2_VL_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=Qwen2VLCausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
-def cce_forward_multimodal(
-    self,
-    input_ids: Optional[torch.LongTensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[list[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    labels: Optional[torch.LongTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    return_dict: Optional[bool] = None,
-    pixel_values: Optional[torch.Tensor] = None,
-    pixel_values_videos: Optional[torch.FloatTensor] = None,
-    image_grid_thw: Optional[torch.LongTensor] = None,
-    video_grid_thw: Optional[torch.LongTensor] = None,
-    rope_deltas: Optional[torch.LongTensor] = None,
-    cache_position: Optional[torch.LongTensor] = None,
-) -> Union[Tuple, Qwen2VLCausalLMOutputWithPast]:
-    r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
-    Returns:
-
-    Example:
-
-    ```python
-    >>> from PIL import Image
-    >>> import requests
-    >>> from transformers import AutoProcessor, Qwen2VLForConditionalGeneration
-
-    >>> model = Qwen2VLForConditionalGeneration.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-    >>> processor = AutoProcessor.from_pretrained("Qwen/Qwen2-VL-7B-Instruct")
-
-    >>> messages = [
-        {
-            "role": "user",
-            "content": [
-                {"type": "image"},
-                {"type": "text", "text": "What is shown in this image?"},
-            ],
-        },
-    ]
-    >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
-    >>> image = Image.open(requests.get(url, stream=True).raw)
-
-    >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
-
-    >>> # Generate
-    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
-    ```"""
-
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-    return_dict = (
-        return_dict if return_dict is not None else self.config.use_return_dict
-    )
-
-    if inputs_embeds is None:
-        inputs_embeds = self.model.embed_tokens(input_ids)
-        if pixel_values is not None:
-            pixel_values = pixel_values.type(self.visual.get_dtype())
-            image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw)
-            n_image_tokens = (input_ids == self.config.image_token_id).sum().item()
-            n_image_features = image_embeds.shape[0]
-            if n_image_tokens != n_image_features:
-                raise ValueError(
-                    f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
-                )
-            image_mask = (
-                (input_ids == self.config.image_token_id)
-                .unsqueeze(-1)
-                .expand_as(inputs_embeds)
-                .to(inputs_embeds.device)
-            )
-            image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)  # type: ignore
-
-        if pixel_values_videos is not None:
-            pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype())
-            video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
-            n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
-            n_video_features = video_embeds.shape[0]
-            if n_video_tokens != n_video_features:
-                raise ValueError(
-                    f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
-                )
-            video_mask = (
-                (input_ids == self.config.video_token_id)
-                .unsqueeze(-1)
-                .expand_as(inputs_embeds)
-                .to(inputs_embeds.device)
-            )
-            video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
-            inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)  # type: ignore
-
-        if attention_mask is not None:
-            attention_mask = attention_mask.to(inputs_embeds.device)
-
-    # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
-    if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
-        # calculate RoPE index once per generation in the pre-fill stage only
-        if (
-            (cache_position is not None and cache_position[0] == 0)
-            or self.rope_deltas is None
-            or (past_key_values is None or past_key_values.get_seq_length() == 0)  # type: ignore
-        ):
-            position_ids, rope_deltas = self.get_rope_index(
-                input_ids, image_grid_thw, video_grid_thw, attention_mask
-            )
-            self.rope_deltas = rope_deltas
-        # then use the prev pre-calculated rope-deltas to get the correct position ids
-        else:
-            batch_size, seq_length, _ = inputs_embeds.shape
-            delta = (
-                cache_position[0] + self.rope_deltas
-                if cache_position is not None
-                else 0
-            )
-            position_ids = torch.arange(seq_length, device=inputs_embeds.device)  # type: ignore
-            position_ids = position_ids.view(1, -1).expand(batch_size, -1)  # type: ignore
-            if cache_position is not None:  # otherwise `deltas` is an int `0`
-                delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)  # type: ignore
-                delta = delta.to(position_ids.device)  # type: ignore
-            position_ids = position_ids.add(delta)  # type: ignore
-            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)  # type: ignore
-
-    outputs = self.model(
-        input_ids=None,
-        position_ids=position_ids,
-        attention_mask=attention_mask,
-        past_key_values=past_key_values,
-        inputs_embeds=inputs_embeds,
-        use_cache=use_cache,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        return_dict=return_dict,
-        cache_position=cache_position,
-    )
-
-    hidden_states = outputs[0]
-    logits = None
-    loss = None
-
-    if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
-        assert labels is not None
-        loss = apply_lce(
-            hidden_states,
-            self.lm_head.weight,
-            labels,
-            _PATCH_OPTS,
-        )
-    else:
-        logits = self.lm_head(hidden_states)
-
-        if labels is not None:
-            # Upcast to float if we need to compute the loss to avoid potential precision issues
-            logits = logits.float()
-            # 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)
-
-    if not return_dict:
-        output = (logits,) + outputs[1:]
-        return (loss,) + output if loss is not None else output
-
-    return Qwen2VLCausalLMOutputWithPast(
-        loss=loss,
-        logits=logits,
-        past_key_values=outputs.past_key_values,
-        hidden_states=outputs.hidden_states,
-        attentions=outputs.attentions,
-        rope_deltas=self.rope_deltas,
-    )
-
-
-def patch_qwen2_vl(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    global _PATCH_OPTS  # pylint: disable=global-statement
-
-    from transformers.models.qwen2_vl import modeling_qwen2_vl
-
-    _PATCH_OPTS = patch_options
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_qwen2_vl.Qwen2VLForConditionalGeneration
-        ), f"Expected a Qwen2VLForConditionalGeneration model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
-
-        return maybe_model
-
-    modeling_qwen2_vl.Qwen2VLForConditionalGeneration.forward = cce_forward_multimodal
-    return None

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

@@ -1,35 +0,0 @@
-"""Qwen3 CCE patch. The model inherits Llama's modeling code and uses the same forward method."""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-)
-
-
-def patch_qwen3(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    from transformers.models.qwen3 import modeling_qwen3
-
-    # Set the _PATCH_OPTS in the llama patch file
-    import axolotl.integrations.cut_cross_entropy.monkeypatch.llama as llama_patch
-
-    llama_patch._PATCH_OPTS = patch_options  # pylint: disable=protected-access
-
-    from axolotl.integrations.cut_cross_entropy.monkeypatch.llama import cce_forward
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_qwen3.Qwen3ForCausalLM
-        ), f"Expected a Qwen3ForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(cce_forward, maybe_model)
-        return maybe_model
-
-    modeling_qwen3.Qwen3ForCausalLM.forward = cce_forward
-    return None

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

@@ -1,194 +0,0 @@
-"""Qwen3 MoE CCE patch. Adapted from transformers v4.51.2"""
-
-# pylint: disable=duplicate-code
-
-from types import MethodType
-from typing import Optional, Union
-
-import torch
-import transformers
-from cut_cross_entropy.transformers.utils import (
-    PatchOptions,
-    TransformersModelT,
-    apply_lce,
-)
-from transformers.modeling_outputs import CausalLMOutputWithPast
-from transformers.models.qwen3_moe.modeling_qwen3_moe import (
-    _CONFIG_FOR_DOC,
-    QWEN3_MOE_INPUTS_DOCSTRING,
-    KwargsForCausalLM,
-    MoeCausalLMOutputWithPast,
-    MoeModelOutputWithPast,
-    load_balancing_loss_func,
-)
-from transformers.processing_utils import Unpack
-from transformers.utils import (
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
-from transformers.utils.deprecation import deprecate_kwarg
-from transformers.utils.generic import can_return_tuple
-
-_PATCH_OPTS: PatchOptions | None = None
-
-
-@can_return_tuple
-@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
-@add_start_docstrings_to_model_forward(QWEN3_MOE_INPUTS_DOCSTRING)
-@replace_return_docstrings(
-    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
-)
-def forward(
-    self,
-    input_ids: Optional[torch.LongTensor] = None,
-    attention_mask: Optional[torch.Tensor] = None,
-    position_ids: Optional[torch.LongTensor] = None,
-    past_key_values: Optional[list[torch.FloatTensor]] = None,
-    inputs_embeds: Optional[torch.FloatTensor] = None,
-    labels: Optional[torch.LongTensor] = None,
-    use_cache: Optional[bool] = None,
-    output_attentions: Optional[bool] = None,
-    output_hidden_states: Optional[bool] = None,
-    output_router_logits: Optional[bool] = None,
-    cache_position: Optional[torch.LongTensor] = None,
-    logits_to_keep: Union[int, torch.Tensor] = 0,
-    **kwargs: Unpack[KwargsForCausalLM],
-) -> MoeCausalLMOutputWithPast:
-    r"""
-        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
-            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
-            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
-            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
-
-        logits_to_keep (`int` or `torch.Tensor`, *optional*):
-            If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all
-            `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that
-            token can save memory, which becomes pretty significant for long sequences or large vocabulary size.
-            If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.
-            This is useful when using packed tensor format (single dimension for batch and sequence length).
-
-    Returns:
-
-    Example:
-
-    ```python
-    >>> from transformers import AutoTokenizer, Qwen3MoeForCausalLM
-
-    >>> model = Qwen3MoeForCausalLM.from_pretrained("Qwen/Qwen3-MoE-15B-A2B")
-    >>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-MoE-15B-A2B")
-
-    >>> prompt = "Hey, are you conscious? Can you talk to me?"
-    >>> inputs = tokenizer(prompt, return_tensors="pt")
-
-    >>> # Generate
-    >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
-    >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
-    "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
-    ```"""
-
-    output_attentions = (
-        output_attentions
-        if output_attentions is not None
-        else self.config.output_attentions
-    )
-    output_router_logits = (
-        output_router_logits
-        if output_router_logits is not None
-        else self.config.output_router_logits
-    )
-
-    output_hidden_states = (
-        output_hidden_states
-        if output_hidden_states is not None
-        else self.config.output_hidden_states
-    )
-
-    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
-    outputs: MoeModelOutputWithPast = self.model(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        position_ids=position_ids,
-        past_key_values=past_key_values,
-        inputs_embeds=inputs_embeds,
-        use_cache=use_cache,
-        output_attentions=output_attentions,
-        output_hidden_states=output_hidden_states,
-        output_router_logits=output_router_logits,
-        cache_position=cache_position,
-        **kwargs,
-    )
-
-    hidden_states = outputs.last_hidden_state
-
-    if hidden_states is None:
-        raise ValueError("hidden_states is None")
-
-    loss = None
-    logits = None
-
-    # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
-    slice_indices = (
-        slice(-logits_to_keep, None)
-        if isinstance(logits_to_keep, int)
-        else logits_to_keep
-    )
-
-    if _PATCH_OPTS is not None and _PATCH_OPTS.use_lce(labels, self.training):
-        assert labels is not None
-        loss = apply_lce(
-            hidden_states[:, slice_indices, :],
-            self.lm_head.weight,
-            labels,
-            _PATCH_OPTS,
-            **kwargs,
-        )
-    else:
-        logits = self.lm_head(hidden_states[:, slice_indices, :])
-
-        if labels is not None:
-            loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
-
-    aux_loss = None
-    if output_router_logits:
-        aux_loss = load_balancing_loss_func(
-            outputs.router_logits,
-            self.num_experts,
-            self.num_experts_per_tok,
-            attention_mask,
-        )
-        if labels is not None:
-            loss += self.router_aux_loss_coef * aux_loss.to(  # type: ignore
-                loss.device  # type: ignore
-            )  # make sure to reside in the same device
-
-    return MoeCausalLMOutputWithPast(
-        loss=loss,
-        aux_loss=aux_loss,  # type: ignore
-        logits=logits,
-        past_key_values=outputs.past_key_values,
-        hidden_states=outputs.hidden_states,
-        attentions=outputs.attentions,
-        router_logits=outputs.router_logits,
-    )
-
-
-def patch_qwen3_moe(
-    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
-    patch_options: PatchOptions,
-) -> TransformersModelT | None:
-    global _PATCH_OPTS  # pylint: disable=global-statement
-
-    from transformers.models.qwen3_moe import modeling_qwen3_moe
-
-    _PATCH_OPTS = patch_options
-
-    if isinstance(maybe_model, transformers.PreTrainedModel):
-        assert isinstance(
-            maybe_model, modeling_qwen3_moe.Qwen3MoeForCausalLM
-        ), f"Expected a Qwen3MoeForCausalLM model. Got {type(maybe_model)}."
-        maybe_model.forward = MethodType(forward, maybe_model)
-
-        return maybe_model
-
-    modeling_qwen3_moe.Qwen3MoeForCausalLM.forward = forward
-    return None

src/axolotl/integrations/kd/chat_template.py

@@ -35,8 +35,6 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
         sequence_len,
         roles_to_train=None,
         train_on_eos=None,
-        train_on_eot=None,
-        eot_tokens=None,
         logprobs_field="logprobs",
         gen_temperature=1.0,
         kd_temperature=1.0,
@@ -52,8 +50,6 @@ class ChatTemplateStrategyWithKD(ChatTemplateStrategy):
             sequence_len,
             roles_to_train=roles_to_train,
             train_on_eos=train_on_eos,
-            train_on_eot=train_on_eot,
-            eot_tokens=eot_tokens,
         )
 
     @property

src/axolotl/prompt_strategies/chat_template.py

@@ -33,7 +33,6 @@ class ChatTemplatePrompter(Prompter):
         message_field_training: Optional[str] = None,
         message_field_training_detail: Optional[str] = None,
         field_messages: str = "messages",
-        field_system: str = "system",
         roles: Optional[Dict[str, List[str]]] = None,
         drop_system_message: bool = False,
     ):
@@ -63,7 +62,6 @@ class ChatTemplatePrompter(Prompter):
         self.message_field_training = message_field_training
         self.message_field_training_detail = message_field_training_detail
         self.field_messages = field_messages
-        self.field_system = field_system
         self.tokenizer = tokenizer
         self.processor: Optional[ProcessorMixin] = processor
         self.chat_template = chat_template
@@ -222,13 +220,10 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
         self,
         prompter: "ChatTemplatePrompter",
         tokenizer,
-        train_on_inputs: bool,
-        sequence_len: int,
-        roles_to_train: Optional[List[str]] = None,
-        train_on_eos: Optional[str] = None,
-        train_on_eot: Optional[str] = None,
-        eot_tokens: Optional[List[str]] = None,
-        split_thinking: Optional[bool] = False,
+        train_on_inputs,
+        sequence_len,
+        roles_to_train=None,
+        train_on_eos=None,
     ):
         super().__init__(prompter, tokenizer, train_on_inputs, sequence_len)
         self.prompter: ChatTemplatePrompter = prompter
@@ -241,88 +236,12 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
             ]
 
         self.train_on_eos = train_on_eos
-        # Backward compatibility, load from train_on_eos
-        self.train_on_eot = train_on_eot if train_on_eot is not None else train_on_eos
-
-        # Default to eos_token if eot_tokens not provided
-        self.eot_tokens = (
-            eot_tokens if eot_tokens is not None else [self.tokenizer.eos_token]
-        )
-        self.split_thinking = split_thinking
-
         self.images = "images"
 
         LOG.debug(
             f"The chat template uses the following properites on the message: {self.prompter.chat_template_msg_variables}"
         )
 
-        self._validate_eot_and_eos_tokens()
-
-    def _validate_eot_and_eos_tokens(self):
-        """
-        - Validates that EOT tokens (or eos_token) are in the chat_template
-        - Checks if EOT tokens are encoded as multiple tokens in the tokenizer.
-        - Checks for potential conflicts between train_on_eos and train_on_eot.
-        """
-        if self.prompter.chat_template is None:
-            # Usually this should not happen
-            LOG.warning(
-                "No chat template provided, skipping EOT and EOS token validation"
-            )
-            return
-
-        # If the EOT token is the same as the EOS token, we need to check differently
-        if len(self.eot_tokens) == 1 and self.eot_tokens[0] == self.tokenizer.eos_token:
-            # Check if the eos_token is in the chat_template or as a variable `eos_token`
-            # Note: we check for `eos_token` in the string, but it could possibly not be a variable
-            if (
-                self.tokenizer.eos_token not in self.prompter.chat_template
-                and "eos_token" not in self.prompter.chat_template
-            ):
-                LOG.warning(
-                    f"EOS token '{self.tokenizer.eos_token}' not found in chat_template. Please check if your template/EOS token is correct."
-                )
-            return
-
-        # Create a new list to store tokens that should be kept
-        valid_eot_tokens = []
-        for token in self.eot_tokens:
-            # Check if EOT token is in the chat_template
-            if token not in self.prompter.chat_template:
-                LOG.warning(f"EOT token '{token}' not found in chat_template.")
-                # Don't add to the valid tokens list
-                continue
-
-            valid_eot_tokens.append(token)
-
-        # Replace the original list with the filtered one
-        self.eot_tokens = valid_eot_tokens
-
-        for token in self.eot_tokens:
-            # If token in template, check if EOT token is in tokenizer and not encoded as multiple tokens
-            token_ids = self.tokenizer.encode(token, add_special_tokens=False)
-            if not token_ids:
-                raise ValueError(
-                    "EOT token encoding failed. Please check if the token is valid and can be encoded."
-                )
-            if token_ids and len(token_ids) > 1:
-                raise ValueError(
-                    f"EOT token '{token}' is encoded as multiple tokens: {token_ids}. Please add it under `tokens: ` in the config "
-                    "or (recommended) override unused added_tokens via `added_tokens_overrides: `."
-                )
-
-        # If eos_token is in eot_tokens and conflict between train_on_eos and train_on_eot, raise an error
-        if (
-            self.tokenizer.eos_token in self.eot_tokens
-            and self.train_on_eos != self.train_on_eot
-        ):
-            raise ValueError(
-                "Conflict between train_on_eos and train_on_eot. eos_token is in eot_tokens and train_on_eos != train_on_eot"
-                f"train_on_eos: {self.train_on_eos}, train_on_eot: {self.train_on_eot}"
-                f"eot_tokens: {self.eot_tokens}"
-                f"eos_token: {self.tokenizer.eos_token}"
-            )
-
     @property
     def supports_batched(self) -> bool:
         # Let calling code know we can handle lists of examples
@@ -366,7 +285,6 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
         if (
             not self.roles_to_train
             and not self.train_on_eos
-            and not self.train_on_eot
             and not self.prompter.message_field_training  # type: ignore
             and not self.prompter.message_field_training_detail  # type: ignore
         ):
@@ -402,7 +320,6 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
         labels = [IGNORE_TOKEN_ID] * len(input_ids)
 
         last_eos_idx = -1
-        last_eot_idx = -1
         for index, turn in enumerate(turns):
             role = turn.get("role")
             content = turn.get("content")
@@ -451,46 +368,25 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
 
                 LOG.debug(f"Labels after processing turn {index}: {labels}")
 
-            # Handle special tokens (EOT and EOS)
-            for token_type, find_func, train_option in [
-                ("EOT", self.find_first_eot_token, self.train_on_eot),
-                ("EOS", self.find_first_eos_token, self.train_on_eos),
-            ]:
-                token_idx = find_func(input_ids, start_idx=turn_end_idx)
-
-                if (
-                    token_idx != -1 and abs(token_idx - turn_end_idx) <= 3
-                ):  # Allow for some template padding
-                    # Update the last token index
-                    if token_type == "EOT":  # nosec B105
-                        last_eot_idx = token_idx
-                    else:
-                        last_eos_idx = token_idx
-
-                    # Set labels if needed for this turn
-                    if train_option == "all" or (
-                        train_option == "turn" and should_train
-                    ):
-                        labels[token_idx] = input_ids[token_idx]
-                        LOG.debug(
-                            f"{token_type} token set for training at index {token_idx}"
-                        )
-                else:
-                    LOG.debug(
-                        f"{token_type} token missing after turn {turn}. {token_type.lower()}_idx: {token_idx}, turn_end_idx: {turn_end_idx}"
-                    )
-
-        # Handle 'last' option for special tokens
-        for token_type, last_idx, train_option in [
-            ("EOT", last_eot_idx, self.train_on_eot),
-            ("EOS", last_eos_idx, self.train_on_eos),
-        ]:
-            if train_option == "last" and last_idx != -1:
-                labels[last_idx] = input_ids[last_idx]
+            # Handle EOS token
+            eos_idx = self.find_first_eos_token(input_ids, start_idx=turn_end_idx)
+            if abs(eos_idx - turn_end_idx) <= 3:  # Allow for some template padding
+                last_eos_idx = eos_idx
+                if self.train_on_eos == "all" or (
+                    self.train_on_eos == "turn" and should_train
+                ):
+                    labels[eos_idx] = input_ids[eos_idx]
+                    LOG.debug(f"EOS token set for training at index {eos_idx}")
+            else:
                 LOG.debug(
-                    f"Last {token_type} token set for training at index {last_idx}"
+                    f"EOS token missing after turn {turn}. eos_idx: {eos_idx}, turn_end_idx: {turn_end_idx}"
                 )
 
+        # Handle 'last' option for train_on_eos
+        if self.train_on_eos == "last" and last_eos_idx != -1:
+            labels[last_eos_idx] = input_ids[last_eos_idx]
+            LOG.debug(f"Last EOS token set for training at index {last_eos_idx}")
+
         LOG.debug(f"Final labels: {labels}")
 
         return {
@@ -506,25 +402,6 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
                 return i
         return -1
 
-    def find_first_eot_token(self, input_ids, start_idx):
-        """Find the first EOT token in the input_ids starting from start_idx."""
-        # Get token IDs for all EOT tokens
-        eot_token_ids = []
-        for token in self.eot_tokens:
-            token_ids = self.tokenizer.encode(token, add_special_tokens=False)
-            if len(token_ids) != 1:
-                raise ValueError(
-                    f"EOT token '{token}' is encoded as multiple tokens: {token_ids}. Please add it under `tokens: ` in the config."
-                )
-
-            eot_token_ids.append(token_ids[0])  # Use the last token ID if multiple
-
-        # Search for any of the EOT token IDs
-        for i in range(start_idx, len(input_ids)):
-            if input_ids[i] in eot_token_ids:
-                return i
-        return -1
-
     def find_turn(self, turns: list[dict], turn_idx: int):
         """
         Locate the starting and ending indices of the specified turn in a conversation.
@@ -611,17 +488,6 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
 
     def get_conversation_thread(self, prompt):
         turns = []
-
-        possible_sys_turn = self.transform_message(
-            prompt[self.prompter.field_messages][0]
-        )
-        if (
-            possible_sys_turn["role"] != "system"
-            and self.prompter.field_system in prompt
-        ):
-            turn = {"role": "system", "content": prompt[self.prompter.field_system]}
-            turns.append(turn)
-
         for message in prompt[self.prompter.field_messages]:
             transformed_message = self.transform_message(message)
 
@@ -657,22 +523,6 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
                 transformed_message["role"], transformed_message["role"]
             )
 
-        # TODO handle reasoning_content with split_thinking
-        # if the role is assistant that we want to use reasoning_content
-        if self.split_thinking and transformed_message["role"] == "assistant":
-            content = transformed_message["content"]
-            pairs = [("<think>", "</think>"), ("<reasoning>", "</reasoning>")]
-            for pair in pairs:
-                if pair[0] in content and pair[1] in content:
-                    start_idx = content.find(pair[0])
-                    end_idx = content.find(pair[1])
-                    thinking_content = content[start_idx + len(pair[0]) : end_idx]
-                    transformed_message["reasoning_content"] = thinking_content.strip()
-                    transformed_message["content"] = content[
-                        end_idx + len(pair[1]) :
-                    ].lstrip()
-                    break
-
         # Determine which keys in the original message were not mapped
         mapped_values = set(self.prompter.message_property_mappings.values())
         remaining_keys = set(message) - mapped_values
@@ -705,9 +555,6 @@ class StrategyLoader:
             "sequence_len": cfg.sequence_len,
             "roles_to_train": ds_cfg.get("roles_to_train", ["assistant"]),
             "train_on_eos": ds_cfg.get("train_on_eos", "turn"),
-            "train_on_eot": ds_cfg.get("train_on_eot", None),
-            "eot_tokens": cfg.get("eot_tokens", None),  # loads from cfg, not ds_cfg
-            "split_thinking": ds_cfg.get("split_thinking", False),
         }
 
     def __call__(

src/axolotl/train.py

@@ -29,7 +29,6 @@ from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFRLTrainerBuil
 from axolotl.core.trainers.mixins.sequence_parallel import (
     SequenceParallelContextManager,
 )
-from axolotl.integrations.base import PluginManager
 from axolotl.logging_config import configure_logging
 from axolotl.utils.dict import DictDefault
 from axolotl.utils.distributed import cleanup_distributed
@@ -534,7 +533,4 @@ def train(
     if not cfg.use_ray:
         cleanup_distributed()
 
-    plugin_manager = PluginManager.get_instance()
-    plugin_manager.post_train(cfg, model)
-
     return model, tokenizer, trainer

src/axolotl/utils/data/rl.py

@@ -204,37 +204,7 @@ def load_prepare_preference_datasets(cfg):
             else:
                 eval_dataset = load_split(cfg.test_datasets, cfg)
         if not eval_dataset:
-            if cfg.val_set_size:
-                # ensure we end up with the same fingerprint by doing rank0 first and being able to cache
-                to_hash_train = (
-                    train_dataset._fingerprint  # pylint: disable=protected-access
-                    + "|"
-                    + str(cfg.val_set_size)
-                    + "|"
-                    + "train"
-                    + "|"
-                    + str(cfg.seed or 42)
-                )
-                to_hash_test = (
-                    train_dataset._fingerprint  # pylint: disable=protected-access
-                    + "|"
-                    + str(cfg.val_set_size)
-                    + "|"
-                    + "test"
-                    + "|"
-                    + str(cfg.seed or 42)
-                )
-                train_fingerprint = md5(to_hash_train)
-                test_fingerprint = md5(to_hash_test)
-                ds_w_test_split = train_dataset.train_test_split(
-                    test_size=cfg.val_set_size,
-                    seed=cfg.seed,
-                    shuffle=False,
-                    train_new_fingerprint=train_fingerprint,
-                    test_new_fingerprint=test_fingerprint,
-                )
-                eval_dataset = ds_w_test_split["test"]
-                train_dataset = ds_w_test_split["train"]
+            eval_dataset = None
 
         if not train_is_preprocessed:
             _save_preprocessed_ds(cfg, cfg.datasets, train_dataset)

src/axolotl/utils/models.py

@@ -53,7 +53,6 @@ from transformers.integrations.deepspeed import (
 )
 
 from axolotl.common.architectures import MOE_ARCH_BLOCK
-from axolotl.integrations.base import PluginManager
 from axolotl.models.mamba import fix_mamba_attn_for_loss
 from axolotl.monkeypatch.multipack import (
     SUPPORTED_MULTIPACK_MODEL_TYPES,
@@ -75,7 +74,6 @@ from axolotl.utils.lora_embeddings import get_linear_embedding_layers
 from axolotl.utils.model_shard_quant import load_sharded_model, load_sharded_model_quant
 
 LOG = logging.getLogger(__name__)
-PLUGIN_MANAGER = PluginManager.get_instance()
 
 MULTIMODAL_AUTO_MODEL_MAPPING = {
     "mllama": MllamaForConditionalGeneration,
@@ -573,8 +571,10 @@ class ModelLoader:
             patch_gemma3conditionalgeneration_forward()
 
         # load any patches from plugins
+        from axolotl.integrations.base import PluginManager
 
-        PLUGIN_MANAGER.pre_model_load(self.cfg)
+        plugin_manager = PluginManager.get_instance()
+        plugin_manager.pre_model_load(self.cfg)
 
         # monkey patch to allow additional Accelerator init kwargs
         if self.cfg.fp8:
@@ -1252,7 +1252,6 @@ class ModelLoader:
 
         try:
             skip_move_to_device = self.build_model(qlora_fsdp)
-            PLUGIN_MANAGER.post_model_build(self.cfg, self.model)
         except Exception as err:  # pylint: disable=broad-exception-caught
             LOG.exception(err)
             raise err
@@ -1332,8 +1331,6 @@ class ModelLoader:
                 before_kbit_train_or_finetune=False,
             )
 
-        PLUGIN_MANAGER.pre_lora_load(self.cfg, self.model)
-
         # ---------------------------------------------------------
         #  load lora or adapter
         # ---------------------------------------------------------
@@ -1395,7 +1392,7 @@ class ModelLoader:
             gc.collect()
             torch.cuda.empty_cache()
 
-        PLUGIN_MANAGER.post_model_load(self.cfg, self.model)
+        # TODO resume_from_checkpoint handling
         return self.model, lora_config
 
 
@@ -1430,13 +1427,9 @@ def load_adapter(model, cfg, adapter, inference=False):
     if hasattr(model, "enable_input_require_grads"):
         model.enable_input_require_grads()
     if adapter in ["lora", "qlora"]:
-        model, lora_config = load_lora(model, cfg, inference=inference)
-        PLUGIN_MANAGER.post_lora_load(cfg, model)
-        return model, lora_config
+        return load_lora(model, cfg, inference=inference)
     if adapter == "llama-adapter":
-        model, lora_config = load_llama_adapter(model, cfg)
-        PLUGIN_MANAGER.post_lora_load(cfg, model)
-        return model, lora_config
+        return load_llama_adapter(model, cfg)
 
     raise NotImplementedError(f"{adapter} peft adapter not available")
 

src/axolotl/utils/schemas/config.py

@@ -309,7 +309,6 @@ class AxolotlInputConfig(
         | Annotated[str, StringConstraints(pattern="^tokenizer_default_fallback_")]
     ) | None = None
     chat_template_jinja: str | None = None
-    eot_tokens: list[str] | None = None
     default_system_message: str | None = None
 
     fix_untrained_tokens: int | list[int] | None = None

src/axolotl/utils/schemas/datasets.py

@@ -50,7 +50,6 @@ class SFTDataset(BaseModel):
     message_property_mappings: dict[str, str] | None = None
     message_field_training: str | None = None
     message_field_training_detail: str | None = None
-    split_thinking: bool | None = None
     logprobs_field: str | None = None
     temperature: float | None = None
     roles_to_train: list[str] | None = None

src/axolotl/utils/schemas/enums.py

@@ -35,7 +35,6 @@ class ChatTemplate(str, Enum):
     jamba = "jamba"  # pylint: disable=invalid-name
     jinja = "jinja"  # pylint: disable=invalid-name
     qwen_25 = "qwen_25"  # pylint: disable=invalid-name
-    qwen3 = "qwen3"  # pylint: disable=invalid-name
     tokenizer_default = "tokenizer_default"  # pylint: disable=invalid-name
     exaone = "exaone"  # pylint: disable=invalid-name
     metharme = "metharme"  # pylint: disable=invalid-name

tests/conftest.py

@@ -79,9 +79,9 @@ def download_smollm2_135m_model():
 
 
 @pytest.fixture(scope="session", autouse=True)
-def download_smollm2_135m_gptq_model():
+def download_llama_68m_random_model():
     # download the model
-    snapshot_download_w_retry("lilmeaty/SmolLM2-135M-Instruct-GPTQ", repo_type="model")
+    snapshot_download_w_retry("JackFram/llama-68m", repo_type="model")
 
 
 @pytest.fixture(scope="session", autouse=True)
@@ -90,12 +90,6 @@ def download_qwen_2_5_half_billion_model():
     snapshot_download_w_retry("Qwen/Qwen2.5-0.5B", repo_type="model")
 
 
-@pytest.fixture(scope="session", autouse=True)
-def download_qwen3_half_billion_model():
-    # download the model
-    snapshot_download_w_retry("Qwen/Qwen3-0.6B", repo_type="model")
-
-
 @pytest.fixture(scope="session", autouse=True)
 def download_tatsu_lab_alpaca_dataset():
     # download the dataset

tests/e2e/integrations/test_hooks.py

@@ -1,184 +0,0 @@
-"""
-e2e tests to make sure all the hooks are fired on the plugin
-"""
-
-import os
-from pathlib import Path
-
-from axolotl.cli.args import TrainerCliArgs
-from axolotl.common.datasets import load_datasets
-from axolotl.integrations.base import BasePlugin
-from axolotl.train import train
-from axolotl.utils.config import normalize_config, prepare_plugins, validate_config
-from axolotl.utils.dict import DictDefault
-
-from ..utils import check_model_output_exists
-
-
-class LogHooksPlugin(BasePlugin):
-    """
-    fixture to capture in a log file each hook that was fired
-    """
-
-    base_dir = Path("/tmp/axolotl-log-hooks")
-
-    def __init__(self):
-        self.base_dir.mkdir(parents=True, exist_ok=True)
-        try:
-            os.remove(self.base_dir.joinpath("plugin_hooks.log"))
-        except FileNotFoundError:
-            pass
-
-    def pre_model_load(self, cfg):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("pre_model_load\n")
-
-    def post_model_build(self, cfg, model):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("post_model_build\n")
-
-    def pre_lora_load(self, cfg, model):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("pre_lora_load\n")
-
-    def post_lora_load(self, cfg, model):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("post_lora_load\n")
-
-    def post_model_load(self, cfg, model):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("post_model_load\n")
-
-    def create_optimizer(self, cfg, trainer):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("create_optimizer\n")
-
-    def get_trainer_cls(self, cfg):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("get_trainer_cls\n")
-
-    def create_lr_scheduler(
-        self, cfg, trainer, optimizer
-    ):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("create_lr_scheduler\n")
-
-    def add_callbacks_pre_trainer(self, cfg, model):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("add_callbacks_pre_trainer\n")
-        return []
-
-    def add_callbacks_post_trainer(
-        self, cfg, trainer
-    ):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("add_callbacks_post_trainer\n")
-        return []
-
-    def post_train(self, cfg, model):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("post_train\n")
-
-    def post_train_unload(self, cfg):  # pylint: disable=unused-argument
-        with open(
-            self.base_dir.joinpath("plugin_hooks.log"), "a", encoding="utf-8"
-        ) as f:
-            f.write("post_train_unload\n")
-
-
-class TestPluginHooks:
-    """
-    e2e tests to make sure all the hooks are fired during the training
-    """
-
-    def test_plugin_hooks(self, temp_dir):
-        # pylint: disable=duplicate-code
-        cfg = DictDefault(
-            {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "plugins": [
-                    "tests.e2e.integrations.test_hooks.LogHooksPlugin",
-                ],
-                "tokenizer_type": "AutoTokenizer",
-                "sequence_len": 1024,
-                "adapter": "lora",
-                "lora_r": 8,
-                "lora_alpha": 16,
-                "lora_dropout": 0.05,
-                "lora_target_linear": True,
-                "val_set_size": 0.02,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
-                "datasets": [
-                    {
-                        "path": "mhenrichsen/alpaca_2k_test",
-                        "type": "alpaca",
-                    },
-                ],
-                "num_epochs": 1,
-                "micro_batch_size": 2,
-                "gradient_accumulation_steps": 1,
-                "output_dir": temp_dir,
-                "learning_rate": 0.00001,
-                "optimizer": "adamw_torch_fused",
-                "lr_scheduler": "cosine",
-                "max_steps": 5,
-                "flash_attention": True,
-                "bf16": "auto",
-            }
-        )
-
-        cfg = validate_config(cfg)
-        prepare_plugins(cfg)
-        normalize_config(cfg)
-        cli_args = TrainerCliArgs()
-        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
-
-        train(cfg=cfg, dataset_meta=dataset_meta)
-        check_model_output_exists(temp_dir, cfg)
-
-        with open(
-            "/tmp/axolotl-log-hooks" + "/plugin_hooks.log", "r", encoding="utf-8"
-        ) as f:
-            file_contents = f.readlines()
-            file_contents = "\n".join(file_contents)
-            assert "pre_model_load" in file_contents
-            assert "post_model_build" in file_contents
-            assert "pre_lora_load" in file_contents
-            assert "post_lora_load" in file_contents
-            assert "post_model_load" in file_contents
-            # assert "create_optimizer" in file_contents  # not implemented yet
-            assert "get_trainer_cls" in file_contents
-            # assert "create_lr_scheduler" in file_contents  # not implemented yet
-            assert "add_callbacks_pre_trainer" in file_contents
-            assert "add_callbacks_post_trainer" in file_contents
-            assert "post_train" in file_contents
-            # assert "post_train_unload" in file_contents  # not called from test train call
-
-        try:
-            os.remove("/tmp/axolotl-log-hooks" + "/plugin_hooks.log")
-        except FileNotFoundError:
-            pass

tests/e2e/multigpu/solo/test_grpo.py

@@ -4,14 +4,11 @@ GRPO test suite
 
 import os
 import random
-import shutil
 import subprocess  # nosec B404
 import sys
-import tempfile
 import time
 from pathlib import Path
 
-import psutil
 import pytest
 import requests
 import yaml
@@ -24,8 +21,8 @@ from tests.e2e.utils import require_vllm
 
 
 def start_vllm(
-    model: str, env: dict, wait: int | None = None, quiet=False, **kwargs
-) -> subprocess.Popen:
+    model: str, env: dict | None = None, wait: int | None = None, quiet=False, **kwargs
+) -> int:
     """
     helper function to start the VLLM server in the background, mostly for testing purposes
     """
@@ -49,41 +46,10 @@ def start_vllm(
     # print out the command to be executed
     print(" ".join(cmd))
 
-    vllm_logging_json = Path(tempfile.mkdtemp()) / "vllm_logging.json"
-    with open(vllm_logging_json, "w", encoding="utf-8") as temp_file:
-        temp_file.write(
-            """{
-  "formatters": {
-    "json": {
-      "class": "pythonjsonlogger.jsonlogger.JsonFormatter"
-    }
-  },
-  "handlers": {
-    "file": {
-      "class": "logging.FileHandler",
-      "formatter": "json",
-      "level": "DEBUG",
-      "filename": "/tmp/vllm.log",
-      "mode": "a"
-    }
-  },
-  "loggers": {
-    "vllm": {
-      "handlers": ["file"],
-      "level": "DEBUG",
-      "propagate": false
-    }
-  },
-  "version": 1
-}"""
-        )
-
-    cmd_env = env.copy()
-    cmd_env.update({"VLLM_LOGGING_CONFIG_PATH": vllm_logging_json})
     # start `trl vllm-serve` command in the background and capture the process id
     process = subprocess.Popen(  # pylint: disable=consider-using-with
         cmd,
-        env=cmd_env,
+        env=env,
         stdout=subprocess.DEVNULL if quiet else subprocess.PIPE,
         stderr=subprocess.DEVNULL if quiet else subprocess.PIPE,
     )  # nosec B603
@@ -92,51 +58,32 @@ def start_vllm(
     print(f"VLLM server process started (PID: {process.pid})")
 
     # wait until the http server is ready, even if it 404s, but timeout after 60 seconds
-    period_seconds = 5
     started = False
     if wait and host and port:
-        for i in range(0, int(wait), period_seconds):
+        for _ in range(int(wait)):
             try:
                 response = requests.get(f"http://{host}:{port}", timeout=1)
-                print(f"{i}: VLLM server (status: {response.status_code})")
                 if int(response.status_code) in [200, 404]:
                     started = True
                     break
-            except requests.exceptions.RequestException as exc:
-                print(f"{i}: VLLM server failed to start: {str(exc)}")
+            except requests.exceptions.RequestException:
+                pass
 
             # also check if the process.pid is still running
             if not process.poll() is None:
                 break
 
-            time.sleep(period_seconds)
+            time.sleep(1)
 
     if wait and not started:
         print(
             f"VLLM server process did not start within {wait} seconds. Please check your server logs."
         )
-        recursive_kill(process)
-        with open("/tmp/vllm.log", "r", encoding="utf-8") as log_file:
-            print(log_file.read())
-        shutil.rmtree("/tmp/vllm.log")
+        process.kill()
         raise RuntimeError(f"VLLM server process did not start within {wait} seconds.")
 
-    # return the process
-    return process
-
-
-def recursive_kill(process: subprocess.Popen):
-    """
-    Recursively kill a process and its children
-    """
-    process = psutil.Process(process.pid)
-    for child in psutil.Process(process.pid).children(recursive=True):
-        child.terminate()
-        child.kill()
-        os.kill(child.pid, 9)
-    process.terminate()
-    process.kill()
-    os.kill(process.pid, 9)
+    # return the process id
+    return process.pid
 
 
 class TestGRPO:
@@ -227,17 +174,16 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
 
         current_env = os.environ.copy()
         env = {
-            "NCCL_P2P_LEVEL": "NVL",
+            "NCCL_P2P_LEVEL": "LOC",
             **current_env,
             "CUDA_VISIBLE_DEVICES": "1",
-            "VLLM_DISABLE_COMPILE_CACHE": "1",
-            # "VLLM_USE_V1": "0",
+            "VLLM_USE_V1": "0",
         }
-        vllm_process = start_vllm(
+        vllm_process_id = start_vllm(
             cfg.base_model,
             env=env,
             quiet=True,
-            wait=300,
+            wait=120,
             gpu_memory_utilization=0.15,
             max_model_len=cfg.vllm.max_model_len,
             enable_prefix_caching=cfg.vllm.enable_prefix_caching,
@@ -256,14 +202,10 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
                     "--main-process-port",
                     f"{get_torch_dist_unique_port()}",
                 ],
-                env={
-                    "NCCL_P2P_LEVEL": "NVL",
-                    "NCCL_DEBUG": "INFO",
-                    **current_env,
-                },
+                env={"NCCL_P2P_LEVEL": "LOC", "NCCL_DEBUG": "INFO", **current_env},
             )
         finally:
-            recursive_kill(vllm_process)
+            os.kill(vllm_process_id, 9)
 
     @pytest.mark.parametrize(
         "num_gpus",
@@ -320,17 +262,16 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
 
         current_env = os.environ.copy()
         env = {
-            "NCCL_P2P_LEVEL": "NVL",  # nccl can be brittle, assume P2P isn't reliable
+            "NCCL_P2P_LEVEL": "LOC",  # nccl can be brittle, assume P2P isn't reliable
             **current_env,
             "CUDA_VISIBLE_DEVICES": "1",
-            "VLLM_DISABLE_COMPILE_CACHE": "1",
-            # "VLLM_USE_V1": "0",
+            "VLLM_USE_V1": "0",
         }
-        vllm_process = start_vllm(
+        vllm_process_id = start_vllm(
             cfg.base_model,
             env=env,
             quiet=True,
-            wait=300,
+            wait=120,
             gpu_memory_utilization=0.15,
             max_model_len=cfg.vllm.max_model_len,
             enable_prefix_caching=cfg.vllm.enable_prefix_caching,
@@ -349,11 +290,7 @@ def oai_gsm8k_transform(cfg, *args, **kwargs):
                     "--main-process-port",
                     f"{get_torch_dist_unique_port()}",
                 ],
-                env={
-                    "NCCL_P2P_LEVEL": "NVL",
-                    "NCCL_DEBUG": "INFO",
-                    **current_env,
-                },
+                env={"NCCL_P2P_LEVEL": "LOC", "NCCL_DEBUG": "INFO", **current_env},
             )
         finally:
-            recursive_kill(vllm_process)
+            os.kill(vllm_process_id, 9)

tests/e2e/patched/test_4d_multipack_llama.py

@@ -28,7 +28,7 @@ class Test4dMultipackLlama(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "base_model": "JackFram/llama-68m",
                 "flash_attention": False,
                 "sdp_attention": True,
                 "sample_packing": True,
@@ -41,9 +41,6 @@ class Test4dMultipackLlama(unittest.TestCase):
                 "lora_target_linear": True,
                 "sequence_len": 1024,
                 "val_set_size": 0.02,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
                 "datasets": [
                     {
                         "path": "mhenrichsen/alpaca_2k_test",
@@ -76,7 +73,7 @@ class Test4dMultipackLlama(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "base_model": "JackFram/llama-68m",
                 "flash_attention": False,
                 "sdp_attention": False,
                 "sample_packing": True,
@@ -89,9 +86,6 @@ class Test4dMultipackLlama(unittest.TestCase):
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
                 "datasets": [
                     {
                         "path": "mhenrichsen/alpaca_2k_test",

tests/e2e/patched/test_fused_llama.py

@@ -32,7 +32,7 @@ class TestFusedLlama(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "base_model": "JackFram/llama-68m",
                 "flash_attention": True,
                 "pad_to_sequence_len": True,
                 "flash_attn_fuse_qkv": True,
@@ -41,7 +41,9 @@ class TestFusedLlama(unittest.TestCase):
                 "sequence_len": 1024,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {

tests/e2e/patched/test_llama_s2_attention.py

@@ -31,8 +31,8 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 16384,
                 "sample_packing": False,
                 "flash_attention": True,
@@ -44,9 +44,7 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
                 "lora_dropout": 0.05,
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "datasets": [
                     {
                         "path": "Yukang/LongAlpaca-12k",
@@ -80,16 +78,14 @@ class TestLlamaShiftedSparseAttention(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 16384,
                 "sample_packing": False,
                 "flash_attention": True,
                 "s2_attention": True,
                 "val_set_size": 0.02,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "datasets": [
                     {
                         "path": "Yukang/LongAlpaca-12k",

tests/e2e/patched/test_lora_llama_multipack.py

@@ -31,8 +31,8 @@ class TestLoraLlama(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "sample_packing": True,
                 "flash_attention": True,
@@ -44,7 +44,9 @@ class TestLoraLlama(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.2,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {
@@ -82,9 +84,9 @@ class TestLoraLlama(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "lilmeaty/SmolLM2-135M-Instruct-GPTQ",
+                "base_model": "TheBlokeAI/jackfram_llama-68m-GPTQ",
                 "model_type": "AutoModelForCausalLM",
-                "tokenizer_type": "AutoTokenizer",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "sample_packing": True,
                 "flash_attention": True,
@@ -98,7 +100,9 @@ class TestLoraLlama(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {

tests/e2e/test_dpo.py

@@ -31,8 +31,8 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -40,9 +40,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "dpo",
                 "datasets": [
                     {
@@ -79,8 +77,8 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -88,9 +86,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "dpo",
                 "rpo_alpha": 0.5,
                 "datasets": [
@@ -128,8 +124,8 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -137,9 +133,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "dpo",
                 "dpo_use_weighting": True,
                 "datasets": [
@@ -178,8 +172,8 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -187,9 +181,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "kto_pair",
                 "datasets": [
                     {
@@ -226,8 +218,8 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -235,9 +227,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "ipo",
                 "datasets": [
                     {
@@ -274,8 +264,8 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -283,9 +273,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "orpo",
                 "orpo_alpha": 0.1,
                 "remove_unused_columns": False,
@@ -326,7 +314,7 @@ class TestDPOLlamaLora(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
+                "base_model": "JackFram/llama-68m",
                 "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
@@ -335,9 +323,7 @@ class TestDPOLlamaLora(unittest.TestCase):
                 "lora_alpha": 32,
                 "lora_dropout": 0.1,
                 "lora_target_linear": True,
-                "special_tokens": {
-                    "pad_token": "<|endoftext|>",
-                },
+                "special_tokens": {},
                 "rl": "kto",
                 "rl_beta": 0.5,
                 "kto_desirable_weight": 1.0,

tests/e2e/test_llama.py

@@ -26,13 +26,15 @@ class TestLlama:
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "trust_remote_code": True,
                 "sequence_len": 512,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {

tests/e2e/test_load_model.py

@@ -26,9 +26,9 @@ class TestLoadModelUtils:
         # load config
         self.cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
-                "tokenizer_config": "HuggingFaceTB/SmolLM2-135M",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
+                "tokenizer_config": "JackFram/llama-68m",
                 "sequence_len": 1024,
                 "load_in_8bit": False,
                 "adapter": "lora",
@@ -38,7 +38,9 @@ class TestLoadModelUtils:
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {

tests/e2e/test_lora_llama.py

@@ -28,8 +28,8 @@ class TestLoraLlama(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -39,7 +39,9 @@ class TestLoraLlama(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {
@@ -48,13 +50,13 @@ class TestLoraLlama(unittest.TestCase):
                     },
                 ],
                 "num_epochs": 1,
-                "micro_batch_size": 2,
+                "micro_batch_size": 8,
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
                 "optimizer": "adamw_torch_fused",
                 "lr_scheduler": "cosine",
-                "max_steps": 5,
+                "max_steps": 20,
             }
         )
 

tests/e2e/test_optimizers.py

@@ -28,9 +28,8 @@ class TestCustomOptimizers(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "model_type": "AutoModelForCausalLM",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -40,7 +39,9 @@ class TestCustomOptimizers(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {
@@ -74,9 +75,8 @@ class TestCustomOptimizers(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "model_type": "AutoModelForCausalLM",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -86,7 +86,9 @@ class TestCustomOptimizers(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {
@@ -120,9 +122,8 @@ class TestCustomOptimizers(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "model_type": "AutoModelForCausalLM",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -132,7 +133,9 @@ class TestCustomOptimizers(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {
@@ -167,7 +170,6 @@ class TestCustomOptimizers(unittest.TestCase):
         cfg = DictDefault(
             {
                 "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "model_type": "AutoModelForCausalLM",
                 "sequence_len": 1024,
                 "val_set_size": 0.01,
                 "special_tokens": {

tests/e2e/test_schedulers.py

@@ -28,8 +28,8 @@ class TestCustomSchedulers(unittest.TestCase):
         # pylint: disable=duplicate-code
         cfg = DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "sequence_len": 1024,
                 "load_in_8bit": True,
                 "adapter": "lora",
@@ -39,7 +39,9 @@ class TestCustomSchedulers(unittest.TestCase):
                 "lora_target_linear": True,
                 "val_set_size": 0.02,
                 "special_tokens": {
-                    "pad_token": "<|endoftext|>",
+                    "unk_token": "<unk>",
+                    "bos_token": "<s>",
+                    "eos_token": "</s>",
                 },
                 "datasets": [
                     {

tests/prompt_strategies/test_chat_templates_advanced.py

@@ -2,8 +2,6 @@
 tests for chat_template prompt strategy
 """
 
-# pylint: disable=too-many-lines
-
 import logging
 from copy import deepcopy
 
@@ -55,6 +53,14 @@ class TestChatTemplateConfigurations:
     Test class for various configurations of ChatTemplateStrategy.
     """
 
+    @staticmethod
+    def find_sublist(full_list, sub_list):
+        token_count = len(sub_list)
+        for index in range(len(full_list) - token_count + 1):
+            if full_list[index : index + token_count] == sub_list:
+                return index
+        return -1
+
     @staticmethod
     def setup_tokenizer(
         tokenizer_name,
@@ -62,7 +68,6 @@ class TestChatTemplateConfigurations:
         chat_template_jinja=None,
         eos_token=None,
         request=None,
-        eot_token=None,
     ) -> tuple[PreTrainedTokenizer, str]:
         """
         Helper function to set up the tokenizer and chat template for the test.
@@ -83,10 +88,6 @@ class TestChatTemplateConfigurations:
                 "CodeLlamaTokenizerFast",
             ):
                 tokenizer.update_post_processor()
-
-        if eot_token:
-            tokenizer.add_special_tokens({"additional_special_tokens": [eot_token]})
-
         return tokenizer, chat_template_jinja
 
     def _should_skip_turn(self, tokenizer, turn, turn_idx, start_idx, end_idx):
@@ -973,311 +974,3 @@ class TestChatTemplateConfigurations:
             raise ValueError(
                 f"Unsupported chat template: {chat_template} with {chat_template_jinja}"
             )
-
-    def test_eot_tokens_conflict_with_eos_token(
-        self,
-        tokenizer,
-        chat_template,
-        chat_template_jinja,
-        eos_token,
-        basic_dataset,  # pylint: disable=unused-argument
-        request,
-    ):
-        """Test that an error is raised when eot_tokens contains eos_token and train_on_eot/train_on_eos conflict"""
-        LOG.info(
-            "Testing conflict between eot_tokens containing eos_token and train_on_eot/train_on_eos mismatch"
-        )
-
-        tokenizer, chat_template_jinja = self.setup_tokenizer(
-            tokenizer, chat_template, chat_template_jinja, eos_token, request
-        )
-
-        # Create a situation where eot_tokens contains eos_token
-        eot_tokens = [
-            tokenizer.eos_token,
-            "[/INST]",
-        ]  # Deliberately including eos_token
-
-        # Create conflicting train_on_eos and train_on_eot settings
-        with pytest.raises(
-            ValueError,
-            match=".*eos_token is in eot_tokens and train_on_eos != train_on_eot.*",
-        ):
-            ChatTemplateStrategy(
-                ChatTemplatePrompter(
-                    tokenizer,
-                    chat_template=get_chat_template(
-                        chat_template, jinja_template=chat_template_jinja
-                    ),
-                    message_property_mappings={"role": "from", "content": "value"},
-                    field_messages="conversations",
-                ),
-                tokenizer=tokenizer,
-                train_on_inputs=False,
-                sequence_len=512,
-                roles_to_train=["assistant"],
-                train_on_eos="none",  # Setting to none
-                train_on_eot="turn",  # Different from train_on_eos
-                eot_tokens=eot_tokens,
-            )
-
-    def test_eot_token_backward_compatibility(
-        self,
-        tokenizer,
-        chat_template,
-        chat_template_jinja,
-        eos_token,
-        basic_dataset,  # pylint: disable=unused-argument
-        request,
-    ):
-        """Test that eot_tokens inherits from eos_token when not specified"""
-        LOG.info("Testing backward compatibility that eot_token inherits eos_token")
-
-        tokenizer, chat_template_jinja = self.setup_tokenizer(
-            tokenizer, chat_template, chat_template_jinja, eos_token, request
-        )
-
-        strategy = ChatTemplateStrategy(
-            ChatTemplatePrompter(
-                tokenizer,
-                chat_template=get_chat_template(
-                    chat_template, jinja_template=chat_template_jinja
-                ),
-                message_property_mappings={"role": "from", "content": "value"},
-                field_messages="conversations",
-            ),
-            tokenizer=tokenizer,
-            train_on_inputs=False,
-            sequence_len=512,
-            roles_to_train=["assistant"],
-            train_on_eos="turn",  # Setting train_on_eos to "turn"
-        )
-
-        # In backward compatibility mode, eot_tokens should be derived from eos_token
-        assert strategy.eot_tokens == [
-            tokenizer.eos_token
-        ], f"Expected eot_tokens to inherit from eos_token, got {strategy.eot_tokens}"
-        assert (
-            strategy.train_on_eot == "turn"
-        ), f"Expected train_on_eot to inherit from train_on_eos, got {strategy.train_on_eot}"
-
-    def test_token_not_in_template(
-        self,
-        tokenizer,
-        chat_template,
-        chat_template_jinja,
-        eos_token,
-        basic_dataset,
-        request,
-    ):
-        """Test runs even when tokens are not found in the template"""
-        LOG.info("Testing runs even when tokens are not found in template")
-
-        tokenizer, chat_template_jinja = self.setup_tokenizer(
-            tokenizer, chat_template, chat_template_jinja, eos_token, request
-        )
-
-        # Create a non-existent token that definitely won't be in the template
-        non_existent_token = "[DEFINITELY_NOT_IN_TEMPLATE]"
-        tokenizer.add_special_tokens(
-            {"additional_special_tokens": [non_existent_token]}
-        )
-
-        strategy = ChatTemplateStrategy(
-            ChatTemplatePrompter(
-                tokenizer,
-                chat_template=get_chat_template(
-                    chat_template, jinja_template=chat_template_jinja
-                ),
-                message_property_mappings={"role": "from", "content": "value"},
-                field_messages="conversations",
-            ),
-            tokenizer=tokenizer,
-            train_on_inputs=False,
-            sequence_len=512,
-            roles_to_train=["assistant"],
-            eot_tokens=[non_existent_token],
-        )
-
-        # Force template check by calling tokenize_prompt
-        strategy.tokenize_prompt(basic_dataset[0])
-
-        # We can also check that a warning was logged, but there's
-        # caplog conflicts when running with other tests
-        # assert any(
-        #     "not found in chat_template" in record.message for record in self._caplog.records
-        # ), "Expected warning about token not found in template was not logged"
-
-    def test_custom_eot_tokens(
-        self,
-        tokenizer,
-        chat_template,
-        chat_template_jinja,
-        eos_token,  # pylint: disable=unused-argument
-        basic_dataset,
-        request,
-    ):
-        """Test with custom EOT tokens to ensure proper masking and training"""
-        LOG.info("Testing with custom EOT tokens")
-
-        tokenizer, chat_template_jinja = self.setup_tokenizer(
-            tokenizer, chat_template, chat_template_jinja, None, request
-        )
-
-        # Add custom EOT tokens to the tokenizer
-        custom_eot = "[EOT]"
-        tokenizer.add_special_tokens({"additional_special_tokens": [custom_eot]})
-
-        # Create a custom chat template that uses our EOT token
-        custom_template = """{% for message in messages %}{% if message['role'] == 'system' %}{{ message['content'] }}{% elif message['role'] == 'user' %}User: {{ message['content'] }}{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}[EOT]{% endif %}{% endfor %}"""
-
-        strategy = ChatTemplateStrategy(
-            ChatTemplatePrompter(
-                tokenizer,
-                chat_template=custom_template,
-                message_property_mappings={"role": "from", "content": "value"},
-                field_messages="conversations",
-            ),
-            tokenizer=tokenizer,
-            train_on_inputs=False,
-            sequence_len=512,
-            roles_to_train=["assistant"],
-            train_on_eot="turn",  # Train on EOT token after each turn
-            eot_tokens=[custom_eot],
-        )
-
-        res = strategy.tokenize_prompt(basic_dataset[0])
-        labels = res["labels"]
-        input_ids = res["input_ids"]
-
-        # Find indices of the EOT token
-        eot_token_id = tokenizer.convert_tokens_to_ids(custom_eot)
-        eot_indices = [
-            i for i, token_id in enumerate(input_ids) if token_id == eot_token_id
-        ]
-
-        assert len(eot_indices) > 0, "Expected at least one EOT token in the input"
-
-        # Verify labeling for EOT tokens based on role
-        turns = strategy.get_conversation_thread(basic_dataset[0])
-        assistant_turn_indices = []
-        non_assistant_turn_indices = []
-
-        for i, turn in enumerate(basic_dataset[0]["conversations"]):
-            start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
-            if start_idx != -1 and end_idx != -1:  # If turn is found
-                if turn["from"] == "assistant":
-                    assistant_turn_indices.append((start_idx, end_idx))
-                else:
-                    non_assistant_turn_indices.append((start_idx, end_idx))
-
-        # Check EOT tokens after assistant turns are labeled
-        for eot_idx in eot_indices:
-            is_after_assistant = any(
-                start_idx <= eot_idx <= end_idx + 1  # +1 to include the EOT token
-                for start_idx, end_idx in assistant_turn_indices
-            )
-
-            if is_after_assistant:
-                assert (
-                    labels[eot_idx] != IGNORE_TOKEN_ID
-                ), f"Expected EOT token after assistant turn at index {eot_idx} to be labeled"
-            else:
-                assert (
-                    labels[eot_idx] == IGNORE_TOKEN_ID
-                ), f"Expected EOT token not after assistant turn at index {eot_idx} to not be labeled"
-
-    def test_multiple_train_on_eot_settings(
-        self,
-        tokenizer,
-        chat_template,
-        chat_template_jinja,
-        eos_token,
-        basic_dataset,
-        request,
-    ):
-        """Test different train_on_eot settings"""
-        LOG.info("Testing different train_on_eot settings")
-
-        tokenizer, chat_template_jinja = self.setup_tokenizer(
-            tokenizer, chat_template, chat_template_jinja, eos_token, request
-        )
-
-        # Create a list to test different train_on_eot settings
-        test_settings = [
-            ("none", lambda idx, is_assistant: False),  # Never train on EOT
-            ("all", lambda idx, is_assistant: True),  # Always train on EOT
-            (
-                "turn",
-                lambda idx, is_assistant: is_assistant,
-            ),  # Train on EOT after assistant turns
-            ("last", lambda idx, is_last: is_last),  # Only train on last EOT
-        ]
-
-        for setting, expected_train_func in test_settings:
-            LOG.info(f"Testing train_on_eot='{setting}'")
-
-            strategy = ChatTemplateStrategy(
-                ChatTemplatePrompter(
-                    tokenizer,
-                    chat_template=get_chat_template(
-                        chat_template, jinja_template=chat_template_jinja
-                    ),
-                    message_property_mappings={"role": "from", "content": "value"},
-                    field_messages="conversations",
-                ),
-                tokenizer=tokenizer,
-                train_on_inputs=False,
-                sequence_len=512,
-                roles_to_train=["assistant"],
-                train_on_eot=setting,
-                eot_tokens=[
-                    tokenizer.eos_token
-                ],  # Use eos_token as the EOT token for simplicity
-            )
-
-            res = strategy.tokenize_prompt(basic_dataset[0])
-            turns = strategy.get_conversation_thread(basic_dataset[0])
-            labels = res["labels"]
-            input_ids = res["input_ids"]
-
-            eos_token_id = tokenizer.eos_token_id
-            eos_indices = [
-                i for i, token_id in enumerate(input_ids) if token_id == eos_token_id
-            ]
-
-            assert (
-                len(eos_indices) > 0
-            ), "Expected at least one EOS/EOT token in the input"
-
-            # Check labeling for each EOS/EOT token
-            for idx, eos_idx in enumerate(eos_indices):
-                # Find which turn this EOS token belongs to
-                preceding_turn = None
-                for i, turn in enumerate(basic_dataset[0]["conversations"]):
-                    start_idx, end_idx = strategy.find_turn(turns=turns, turn_idx=i)
-                    if (
-                        start_idx != -1
-                        and end_idx != -1
-                        and start_idx <= eos_idx <= end_idx + 1
-                    ):
-                        preceding_turn = turn
-                        break
-
-                is_assistant = (
-                    preceding_turn is not None and preceding_turn["from"] == "assistant"
-                )
-                is_last = idx == len(eos_indices) - 1
-
-                expected_label = not expected_train_func(
-                    idx, is_assistant if setting != "last" else is_last
-                )
-
-                if expected_label:
-                    assert (
-                        labels[eos_idx] == IGNORE_TOKEN_ID
-                    ), f"Expected EOT token at index {eos_idx} to not be labeled with train_on_eot='{setting}'"
-                else:
-                    assert (
-                        labels[eos_idx] != IGNORE_TOKEN_ID
-                    ), f"Expected EOT token at index {eos_idx} to be labeled with train_on_eot='{setting}'"

tests/prompt_strategies/test_chat_templates_thinking.py

@@ -1,118 +0,0 @@
-"""
-Tests for splitting reasoning/thinking from content into separate field
-"""
-
-import logging
-
-import pytest
-from datasets import Dataset
-from transformers import AutoTokenizer
-
-from axolotl.prompt_strategies.chat_template import (
-    load,
-)
-from axolotl.utils.dict import DictDefault
-
-from tests.hf_offline_utils import enable_hf_offline
-
-logging.basicConfig(level=logging.DEBUG)
-LOG = logging.getLogger("axolotl")
-
-
-@pytest.fixture(name="messages_w_reasoning")
-def messages_w_reasoning_fixture():
-    return Dataset.from_list(
-        [
-            {
-                "messages": [
-                    {
-                        "role": "user",
-                        "content": "hello",
-                    },
-                    {
-                        "role": "assistant",
-                        "content": "<think>lorem</think>\nwelcome",
-                    },
-                ]
-            }
-        ]
-    )
-
-
-@pytest.fixture(name="qwen3_tokenizer")
-@enable_hf_offline
-def qwen3_tokenizer_fixture(
-    download_qwen3_half_billion_model,
-):  # pylint: disable=unused-argument
-    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
-
-    return tokenizer
-
-
-class TestSplitThinking:
-    """
-    test class to make sure datasets with reasoning content conforms to the chat_template strategy
-    """
-
-    def test_splits_think(self, messages_w_reasoning, qwen3_tokenizer):
-        # pylint: disable=duplicate-code
-        strategy = load(
-            qwen3_tokenizer,
-            DictDefault(
-                {
-                    "train_on_inputs": False,
-                    "sequence_len": 512,
-                }
-            ),
-            DictDefault(
-                {
-                    "chat_template": "qwen3",
-                    "message_field_role": "role",
-                    "message_field_content": "content",
-                    "message_property_mappings": {
-                        "role": "role",
-                        "content": "content",
-                    },
-                    "roles": {
-                        "user": ["user"],
-                        "assistant": ["assistant"],
-                        "system": ["system"],
-                    },
-                    "field_messages": "messages",
-                    "split_thinking": True,
-                }
-            ),
-        )
-        transformed_prompt = strategy.get_conversation_thread(messages_w_reasoning[0])
-        assert transformed_prompt[0]["role"] == "user"
-        assert transformed_prompt[1]["role"] == "assistant"
-        assert transformed_prompt[1]["reasoning_content"] == "lorem"
-        assert transformed_prompt[1]["content"] == "welcome"
-
-        res = strategy.tokenize_prompt(messages_w_reasoning[0])
-        input_ids = res["input_ids"]
-        # fmt: off
-        expected_input_ids = [
-            151644,  # im_start
-            872,  # user
-            198,  # \n
-            14990,  # hello
-            151645,  # im_end
-            198,  # \n
-            151644,  # im_start
-            77091,  # assistant
-            198,  # \n
-            151667,  # think
-            198,  # \n
-            385, 1826,  # lorem
-            198,  # \n
-            151668,  # /think
-            271,  # \n
-            34084,  # welcome
-            151645,  # im_end
-            198,  # \n
-        ]
-        # fmt: on
-        assert (
-            input_ids == expected_input_ids
-        ), f"Input IDs mismatch: {input_ids} != {expected_input_ids}"

tests/test_normalize_config.py

@@ -17,9 +17,9 @@ class NormalizeConfigTestCase(unittest.TestCase):
     def _get_base_cfg(self):
         return DictDefault(
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "base_model_config": "HuggingFaceTB/SmolLM2-135M",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "base_model_config": "JackFram/llama-68m",
+                "tokenizer_type": "LlamaTokenizer",
                 "num_epochs": 1,
                 "micro_batch_size": 1,
                 "gradient_accumulation_steps": 1,

tests/utils/test_models.py

@@ -18,9 +18,9 @@ class TestModelsUtils:
         # load config
         self.cfg = DictDefault(  # pylint: disable=attribute-defined-outside-init
             {
-                "base_model": "HuggingFaceTB/SmolLM2-135M",
-                "model_type": "AutoModelForCausalLM",
-                "tokenizer_type": "AutoTokenizer",
+                "base_model": "JackFram/llama-68m",
+                "model_type": "LlamaForCausalLM",
+                "tokenizer_type": "LlamaTokenizer",
                 "load_in_8bit": True,
                 "load_in_4bit": False,
                 "adapter": "lora",
@@ -65,7 +65,7 @@ class TestModelsUtils:
                 "s2_attention": True,
                 "sample_packing": True,
                 "base_model": "",
-                "model_type": "AutoModelForCausalLM",
+                "model_type": "LlamaForCausalLM",
             }
         )