fix dumb mistakes

* add override of upstream fix

* override batch loss metrics for CPO/Simpo as well

.github/workflows/docs.yml

@@ -23,6 +23,7 @@ jobs:
         - name: Install dependencies
           run: |
             python3 -m pip install jupyter quartodoc
+            python3 -m pip install -e . --no-deps
         - name: Build autodoc
           run: quartodoc build
         - name: Publish to GitHub Pages (and render)

_quarto.yml

@@ -133,6 +133,7 @@ quartodoc:
         - utils.schemas.datasets
         - utils.schemas.peft
         - utils.schemas.trl
+        - utils.schemas.multimodal
         - utils.schemas.integrations
         - utils.schemas.enums
         - utils.schemas.utils

docs/config.qmd

@@ -466,6 +466,7 @@ auto_find_batch_size: # Optional[bool]
 
 eval_table_size: # Approximate number of predictions sent to wandb depending on batch size. Enabled above 0. Default is 0
 eval_max_new_tokens: # Total number of tokens generated for predictions sent to wandb. Default is 128
+do_causal_lm_eval: # Whether to run causal language model evaluation for metrics in `eval_causal_lm_metrics`.
 eval_causal_lm_metrics: # HF evaluate metrics used during evaluation. Default is ["sacrebleu", "comet", "ter", "chrf", "perplexity"]
 
 profiler_steps: # enable the pytorch profiler to capture the first N steps of training to the output_dir.
@@ -506,36 +507,58 @@ lr_div_factor: # Learning rate div factor
 
 # Specify optimizer
 # Valid values are driven by the Transformers OptimizerNames class, see:
-# https://github.com/huggingface/transformers/blob/95b374952dc27d8511541d6f5a4e22c9ec11fb24/src/transformers/training_args.py#L134
+# https://github.com/huggingface/transformers/blob/cbf924b76c03828101a34069a96d209314114fd5/src/transformers/training_args.py#L144-L189
 #
 # Note that not all optimizers may be available in your environment, ex: 'adamw_anyprecision' is part of
 # torchdistx, 'adamw_bnb_8bit' is part of bnb.optim.Adam8bit, etc. When in doubt, it is recommended to start with the optimizer used
 # in the examples/ for your model and fine-tuning use case.
 #
 # Valid values for 'optimizer' include:
-# - adamw_hf
 # - adamw_torch
 # - adamw_torch_fused
 # - adamw_torch_xla
+# - adamw_torch_npu_fused
 # - adamw_apex_fused
-# - adopt_adamw (an EXPERIMENTAL optimizer, only for torch version >= 2.5.1)
+# - adopt_adamw  (an EXPERIMENTAL optimizer, only for torch version >= 2.5.1)
 # - adafactor
 # - adamw_anyprecision
+# - adamw_torch_4bit
+# - ademamix
 # - sgd
 # - adagrad
 # - adamw_bnb_8bit
+# - adamw_8bit   # alias for adamw_bnb_8bit
+# - ademamix_8bit
 # - lion_8bit
 # - lion_32bit
 # - paged_adamw_32bit
 # - paged_adamw_8bit
+# - paged_ademamix_32bit
+# - paged_ademamix_8bit
 # - paged_lion_32bit
 # - paged_lion_8bit
+# - rmsprop
+# - rmsprop_bnb
+# - rmsprop_bnb_8bit
+# - rmsprop_bnb_32bit
 # - galore_adamw
 # - galore_adamw_8bit
 # - galore_adafactor
 # - galore_adamw_layerwise
 # - galore_adamw_8bit_layerwise
 # - galore_adafactor_layerwise
+# - lomo
+# - adalomo
+# - grokadamw
+# - schedule_free_adamw
+# - schedule_free_sgd
+# - apollo_adamw
+# - apollo_adamw_layerwise
+#
+# Additional custom optimizers include:
+# - optimi_adamw
+# - ao_adamw_8bit
+# - ao_adamw_fp8
 optimizer:
 # Dictionary of arguments to pass to the optimizer
 optim_args:
@@ -587,6 +610,14 @@ resume_from_checkpoint:
 # Be careful with this being turned on between different models.
 auto_resume_from_checkpoints: false
 
+## Multimodal section
+# int | tuple[int, int] | None . Size to resize images to, width x height.
+# Will read from model/processor config if not set.
+image_size:
+# str. Algorithm to use for image resizing. "bilinear", "bicubic", "lanczos". Default is "bilinear".
+image_resize_algorithm: 'bilinear'
+## End of multimodal section
+
 # Don't mess with this, it's here for accelerate and torchrun
 local_rank:
 

docs/docker.qmd

@@ -103,8 +103,7 @@ This uses the same tags as the [`main` image](#sec-main-tags).
 
 - `JUPYTER_DISABLE`: Disable Jupyter lab.
 - `JUPYTER_PASSWORD`: Set a password for the Jupyter lab.
-- `PUBLIC_KEY`: Add a public key for the SSH service.
+- `PUBLIC_KEY` / `SSH_KEY`: Add a public key for the SSH service.
-- `SSH_KEY`: Add a private key for the SSH service.
 
 #### Volume mounts
 

docs/multimodal.qmd

@@ -1,28 +1,171 @@
-# MultiModal / Vision Language Models (BETA)
+---
+title: MultiModal / Vision Language Models (BETA)
+format:
+  html:
+    toc: true
+    toc-depth: 3
+---
 
-### Supported Models
+## Supported Models
 
-- Mllama, i.e. llama with vision models
+- [Mllama](#sec-mllama)
+- [Pixtral](#sec-pixtral)
+- [Llava-1.5](#sec-llava-15)
+- [Mistral-Small-3.1](#sec-mistral-small-31)
+- [Gemma-3](#sec-gemma-3)
+- [Qwen2-VL](#sec-qwen2-vl)
+- [Qwen2.5-VL](#sec-qwen25-vl)
 
-### Usage
+## Usage
 
-Currently multimodal support is limited and doesn't have full feature parity. To finetune a multimodal Llama w/ LoRA,
+Multimodal support is limited and doesn't have full feature parity.
-you'll need to use the following in YAML in combination with the rest of the required hyperparams.
+
+Here are the hyperparams you'll need to use to finetune a multimodal model.
 
 ```yaml
-base_model: alpindale/Llama-3.2-11B-Vision-Instruct
 processor_type: AutoProcessor
-skip_prepare_dataset: true
 
-chat_template: llama3_2_vision
+skip_prepare_dataset: true
+remove_unused_columns: false  # leave columns in place as they are needed to handle image embeddings during training
+sample_packing: false  # not yet supported with multimodal
+
+chat_template:  # see in next section
+
+# example dataset
 datasets:
   - path: HuggingFaceH4/llava-instruct-mix-vsft
     type: chat_template
     split: train[:1%]
     field_messages: messages
-remove_unused_columns: false
-sample_packing: false
 
-# only finetune the Language model, leave the vision model and vision tower frozen
+# (optional) if doing lora, only finetune the Language model,
+# leave the vision model and vision tower frozen
+# load_in_8bit: true
+adapter: lora
 lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+# (optional) if you want to resize images to a set size
+image_size: 512
+image_resize_algorithm: bilinear
+```
+
+Please see [examples](https://github.com/axolotl-ai/axolotl/tree/main/examples) folder for full configs.
+
+::: {.callout-warning}
+Some of our chat_templates have been extended to support broader dataset types. This should not break any existing configs.
+:::
+
+### Mllama {#sec-mllama}
+
+```yaml
+base_model: meta-llama/Llama-3.2-11B-Vision-Instruct
+
+chat_template: llama3_2_vision
+```
+
+### Pixtral {#sec-pixtral}
+
+```yaml
+base_model: mistralai/Pixtral-12B-2409
+
+chat_template: pixtral
+```
+
+### Llava-1.5 {#sec-llava-15}
+
+```yaml
+base_model: llava-hf/llava-1.5-7b-hf
+
+chat_template: llava
+```
+
+### Mistral-Small-3.1 {#sec-mistral-small-31}
+
+```yaml
+base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
+
+chat_template: mistral_v7_tekken
+```
+
+### Gemma-3 {#sec-gemma-3}
+
+::: {.callout-tip}
+The Gemma3-1B model is a text-only model, so please train as regular text model.
+:::
+
+For multi-modal 4B/12B/27B models, use the following config:
+
+```yaml
+base_model: google/gemma-3-4b-it
+
+chat_template: gemma3
+```
+
+### Qwen2-VL {#sec-qwen2-vl}
+
+```yaml
+base_model: Qwen/Qwen2-VL-7B-Instruct
+
+chat_template: qwen2_vl
+```
+
+### Qwen2.5-VL {#sec-qwen25-vl}
+
+```yaml
+base_model: Qwen/Qwen2.5-VL-7B-Instruct
+
+chat_template: qwen2_vl  # same as qwen2-vl
+```
+
+## Dataset Format
+
+For multi-modal datasets, we adopt an extended `chat_template` format similar to OpenAI's Message format.
+
+- A message is a list of `role` and `content`.
+- `role` can be `system`, `user`, `assistant`, etc.
+- `content` is a list of `type` and (`text` or `image` or `path` or `url` or `base64`).
+
+::: {.callout-note}
+For backwards compatibility:
+
+- If the dataset has a `images` or `image` column of `list[Image]`, it will be appended to the first `content` list as `{"type": "image", "image": ...}`. However, if the content already has a `{"type": "image"}` but no `image` key, it will be set the `image` key.
+- If `content` is a string, it will be converted to a list with `type` as `text`.
+:::
+
+::: {.callout-tip}
+For image loading, you can use the following keys within `content` alongside `"type": "image"`:
+
+- `"path": "/path/to/image.jpg"`
+- `"url": "https://example.com/image.jpg"`
+- `"base64": "..."`
+- `"image": PIL.Image`
+:::
+
+Here is an example of a multi-modal dataset:
+```json
+[
+  {
+    "messages": [
+        {
+            "role": "system",
+            "content": [
+              {"type": "text", "text": "You are a helpful assistant."}
+              ]
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "image", "image": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg"},
+                {"type": "text", "text": "Describe this image in detail."}
+            ]
+        },
+        {
+            "role": "assistant",
+            "content": [
+              {"type": "text", "text": "The image is a bee."}
+            ]
+        }
+    ]
+  }
+]
 ```

examples/cohere/command-r-7b-qlora.yml

@@ -0,0 +1,71 @@
+base_model: CohereForAI/c4ai-command-r7b-12-2024
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+
+load_in_8bit: false
+load_in_4bit: true
+strict: false
+
+# huggingface repo
+chat_template: cohere
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: qlora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: false
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 4
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+
+warmup_ratio: 0.1
+evals_per_epoch:
+eval_table_size:
+eval_max_new_tokens: 128
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:

examples/gemma3/gemma-3-4b-lora.yml

@@ -0,0 +1,63 @@
+base_model: google/gemma-3-4b-it
+processor_type: AutoProcessor
+strict: false
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: gemma3
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+    field_messages: messages
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+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: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+local_rank:
+logging_steps: 1
+flash_attention: true
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:

examples/gemma3/qlora.yml

@@ -0,0 +1,74 @@
+base_model: google/gemma-3-1b-it
+# optionally might have model_type or tokenizer_type
+model_type: AutoModelForCausalLM
+tokenizer_type: AutoTokenizer
+# 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
+
+# huggingface repo
+chat_template: gemma3_text
+datasets:
+  - path: cgato/SlimOrcaDedupCleaned
+    type: chat_template
+    field_messages: conversations
+    message_property_mappings:
+      role: from
+      content: value
+
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: qlora
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_linear: true
+
+sequence_len: 2048
+sample_packing: true
+eval_sample_packing: false
+pad_to_sequence_len: true
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 4
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: auto
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+early_stopping_patience:
+resume_from_checkpoint:
+local_rank:
+logging_steps: 1
+xformers_attention:
+flash_attention: true
+
+warmup_ratio: 0.1
+evals_per_epoch:
+eval_table_size:
+eval_max_new_tokens: 128
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:

examples/llava/lora-7b.yaml

@@ -0,0 +1,63 @@
+base_model: llava-hf/llava-1.5-7b-hf
+processor_type: AutoProcessor
+strict: false
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: llava
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+    field_messages: messages
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+local_rank:
+logging_steps: 1
+flash_attention: true
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:

examples/mistral/mistral-small-3.1-24B-lora.yml

@@ -0,0 +1,66 @@
+base_model: mistralai/Mistral-Small-3.1-24B-Instruct-2503
+processor_type: AutoProcessor
+strict: false
+
+load_in_8bit: true
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: mistral_v7_tekken
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+    field_messages: messages
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.01
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 2048
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 1
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+local_rank:
+logging_steps: 1
+flash_attention: false # PixtralVisionModel does not support Flash Attention 2.0 yet.
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:

examples/pixtral/lora-12b.yml

@@ -0,0 +1,65 @@
+base_model: mistral-community/pixtral-12b
+processor_type: AutoProcessor
+strict: false
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: pixtral
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+    field_messages: messages
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'language_model.model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+local_rank:
+logging_steps: 1
+flash_attention: false # PixtralVisionModel does not support Flash Attention 2.0 yet
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:
+special_tokens:
+  pad_token: <pad>

examples/qwen2-vl/lora-7b.yaml

@@ -0,0 +1,63 @@
+base_model: Qwen/Qwen2-VL-7B-Instruct
+processor_type: AutoProcessor
+strict: false
+
+# these 3 lines are needed for now to handle vision chat templates w images
+skip_prepare_dataset: true
+remove_unused_columns: false
+sample_packing: false
+
+chat_template: qwen2_vl
+datasets:
+  - path: HuggingFaceH4/llava-instruct-mix-vsft
+    type: chat_template
+    split: train[:1%]
+    field_messages: messages
+dataset_prepared_path: last_run_prepared
+val_set_size: 0.0
+output_dir: ./outputs/out
+
+adapter: lora
+lora_model_dir:
+
+sequence_len: 8192
+pad_to_sequence_len: false
+
+lora_r: 32
+lora_alpha: 16
+lora_dropout: 0.05
+lora_target_modules: 'model.layers.[\d]+.(mlp|cross_attn|self_attn).(up|down|gate|q|k|v|o)_proj'
+
+wandb_project:
+wandb_entity:
+wandb_watch:
+wandb_name:
+wandb_log_model:
+
+gradient_accumulation_steps: 4
+micro_batch_size: 1
+num_epochs: 1
+optimizer: adamw_bnb_8bit
+lr_scheduler: cosine
+learning_rate: 0.0002
+
+train_on_inputs: false
+group_by_length: false
+bf16: true
+fp16:
+tf32: true
+
+gradient_checkpointing: true
+local_rank:
+logging_steps: 1
+flash_attention: true
+eager_attention:
+
+warmup_ratio: 0.1
+evals_per_epoch: 1
+saves_per_epoch: 1
+debug:
+deepspeed:
+weight_decay: 0.0
+fsdp:
+fsdp_config:

requirements.txt

@@ -12,7 +12,7 @@ liger-kernel==0.5.3
 packaging==23.2
 
 peft==0.15.0
-transformers==4.49.0
+transformers==4.50.0
 tokenizers>=0.21.1
 accelerate==1.5.2
 datasets==3.4.1

requirements_env.txt

@@ -1,315 +0,0 @@
-accelerate==0.34.1
-addict==2.4.0
-aiofiles==23.2.1
-aiohttp==3.9.0
-aiosignal==1.3.1
-aiostream==0.5.2
-alembic==1.13.1
-annotated-types==0.6.0
-annoy==1.17.3
-ansible==6.7.0
-ansible-core==2.13.13
-ansible-vault==2.1.0
-anyio==3.7.1
-appdirs==1.4.4
-art==6.0
-asgiref==3.7.2
-async-timeout==4.0.2
-attrdict==2.0.1
-attrs==22.2.0
-awscli==1.32.75
--e git+ssh://git@github.com/OpenAccess-AI-Collective/axolotl.git@6e354682e3c1735d3f7fb9e362280c38e922260f#egg=axolotl
-backoff==2.2.1
-base58==2.1.1
-beartype==0.17.2
-bitnet==0.2.1
-bitsandbytes==0.42.0
-bittensor==6.7.0
-black==23.7.0
-blinker==1.7.0
-boto3==1.34.75
-botocore==1.34.75
-cachetools==5.3.3
-cachy==0.1.1
-certifi==2023.7.22
-cffi==1.16.0
-cfgv==3.3.1
-chai-guanaco==1.2.4
-charset-normalizer==3.2.0
-cleo==0.6.8
-click==8.1.7
-cloudpickle==2.0.0
-cohere==4.11.2
-colorama==0.4.4
-coloredlogs==15.0.1
-CoLT5-attention==0.10.20
-contextlib2==21.6.0
-contourpy==1.2.0
-cryptography==41.0.3
-cycler==0.12.1
-cytoolz==0.12.3
-databricks-cli==0.18.0
-dataclasses-json==0.5.7
-datasets==2.11.0
-ddt==1.6.0
-decorator==5.1.1
-deepspeed==0.15.0
-# Editable Git install with no remote (dialogpt==0.1)
--e /Users/wing/Projects/ml/dialogpt/src
-dill==0.3.6
-distlib==0.3.6
-docker==7.0.0
-docker-pycreds==0.4.0
-docstring-parser==0.15
-docutils==0.16
-ecdsa==0.18.0
-einops==0.7.0
-einops-exts==0.0.4
-einx==0.1.3
-entrypoints==0.4
-eth-hash==0.6.0
-eth-keys==0.5.0
-eth-typing==4.0.0
-eth-utils==2.3.1
-evaluate==0.4.0
-exceptiongroup==1.1.1
-fastapi==0.109.2
-fastcore==1.5.29
-ffmpy==0.4.0
-filelock==3.12.2
--e git+https://github.com/NousResearch/finetuning-subnet.git@24e9407d6b4430a7ca39d344692f89ce5a97d27e#egg=finetuning_subnet
-fire==0.5.0
-first==2.0.2
-flake8==7.0.0
-Flask==3.0.1
-fonttools==4.47.2
-frozendict==2.4.1
-frozenlist==1.3.3
-fschat @ git+https://github.com/lm-sys/FastChat.git@27a05b04a35510afb1d767ae7e5990cbd278f8fe
-fsspec==2023.6.0
-fuzzywuzzy==0.18.0
-gitdb==4.0.10
-GitPython==3.1.31
-google-pasta==0.2.0
-gradio==4.42.0
-gradio_client==1.3.0
-greenlet==2.0.2
-grpclib==0.4.7
-gunicorn==21.2.0
-h11==0.14.0
-h2==4.1.0
-hpack==4.0.0
-httpcore==0.17.3
-httpx==0.24.1
-huggingface-hub==0.23.4
-humanfriendly==10.0
-hyperframe==6.0.1
-identify==2.5.24
-idna==3.4
-immutables==0.20
-importlib-metadata==6.7.0
-importlib-resources==6.1.1
-inflection==0.5.1
-iniconfig==2.0.0
-itsdangerous==2.1.2
-Jinja2==3.1.2
-jmespath==1.0.1
-joblib==1.3.2
-jsonlines==3.1.0
-jsonschema==2.6.0
-kiwisolver==1.4.5
-langchain==0.0.144
-Levenshtein==0.24.0
-libcst==1.1.0
-liger-kernel==0.0.0
-lion-pytorch==0.1.2
-llama-cpp-python==0.1.36
-llvmlite==0.40.1
-local-attention==1.9.0
-loguru==0.7.0
-Mako==1.3.2
-Markdown==3.5.2
-markdown-it-py==3.0.0
-markdown2==2.4.10
-MarkupSafe==2.1.2
-marshmallow==3.19.0
-marshmallow-enum==1.5.1
-matplotlib==3.8.2
-mccabe==0.7.0
-mdurl==0.1.2
-MEGABYTE-pytorch==0.0.7
--e git+https://github.com/cg123/mergekit.git@53c5f414774a0558b8d84858fb6374bc93a8f1c1#egg=mergekit
-mlflow==2.10.0
-modal==0.62.77
-more-itertools==10.2.0
-mpmath==1.2.1
-msgpack==1.0.7
-msgpack-numpy-opentensor==0.5.0
-multidict==6.0.4
-multiprocess==0.70.14
-munch==2.5.0
-mypy==1.3.0
-mypy-extensions==1.0.0
-nest-asyncio==1.6.0
-netaddr==0.10.1
-networkx==3.0rc1
-nh3==0.2.14
-nodeenv==1.8.0
-nomic==2.0.2
-numba==0.57.1
-numexpr==2.8.4
-numpy==1.24.4
-oauthlib==3.2.2
-openai==0.27.4
-openapi==1.1.0
-openapi-schema-pydantic==1.2.4
-optimum==1.8.6
-orjson==3.10.7
-packaging==23.1
-pandas==2.0.0
-parameterized==0.9.0
-password-strength==0.0.3.post2
-pastel==0.1.1
-pathos==0.3.0
-pathspec==0.11.1
-pathtools==0.1.2
-peft==0.11.1
-pendulum==3.0.0
-Pillow==9.5.0
-pip-tools==1.11.0
-platformdirs==3.2.0
-pluggy==1.4.0
-poetry==0.7.1
-pox==0.3.2
-ppft==1.7.6.6
-pre-commit==3.3.2
-prettytable==3.10.0
-prompt-toolkit==3.0.39
-protobuf==3.20.2
-protobuf3-to-dict==0.1.5
-psutil==5.9.5
-psycopg==3.1.18
-PuLP==2.8.0
-py==1.11.0
-py-bip39-bindings==0.1.11
-py-cpuinfo==9.0.0
-py-ed25519-zebra-bindings==1.0.1
-py-sr25519-bindings==0.2.0
-pyarrow==11.0.0
-pyasn1==0.6.0
-pycodestyle==2.11.1
-pycparser==2.21
-pycryptodome==3.20.0
-pydantic==2.5.3
-pydantic_core==2.14.6
-pydub==0.25.1
-pyfiglet==0.8.post1
-pyflakes==3.2.0
-Pygments==2.15.1
-PyJWT==2.8.0
-pylev==1.4.0
-PyNaCl==1.5.0
-pynvml==11.5.0
-pyparsing==2.4.7
-pyrsistent==0.14.11
-pytest==8.0.2
-pytest-asyncio==0.23.4
-python-dateutil==2.8.2
-python-dotenv==1.0.1
-python-Levenshtein==0.24.0
-python-multipart==0.0.9
-pytz==2023.3
-PyYAML==6.0.1
-querystring-parser==1.2.4
-rapidfuzz==3.6.1
-regex==2023.6.3
-requests==2.31.0
-requests-toolbelt==0.8.0
-resolvelib==0.8.1
-responses==0.18.0
-retry==0.9.2
-rich==13.7.0
-rsa==4.7.2
-ruff==0.6.3
-s3transfer==0.10.1
-safetensors==0.4.5
-sagemaker==2.148.0
-scalecodec==1.2.7
-schedulefree==1.2.1
-schema==0.7.5
-scikit-learn==1.4.0
-scipy==1.9.3
-seaborn==0.13.2
-semantic-version==2.10.0
-sentencepiece==0.2.0
-sentry-sdk==1.19.1
-setproctitle==1.3.2
-shellingham==1.5.4
-shortuuid==1.0.11
-shtab==1.6.5
-sigtools==4.0.1
-six==1.16.0
-skypilot==0.4.1
-smdebug-rulesconfig==1.0.1
-smmap==5.0.0
-sniffio==1.3.0
-SQLAlchemy==1.4.47
-sqlparse==0.4.4
-starlette==0.36.3
-substrate-interface==1.5.2
-svgwrite==1.4.3
-sympy==1.11.1
-synchronicity==0.6.7
-tabulate==0.9.0
-tblib==1.7.0
-tenacity==8.2.2
-tensor-parallel==2.0.0
-termcolor==2.2.0
-text2art==0.2.0
-threadpoolctl==3.2.0
-tiktoken==0.6.0
-time-machine==2.14.1
-timm==0.9.16
-tokenizers==0.19.1
-tokenmonster==1.1.12
-toml==0.9.6
-tomli==2.0.1
-tomlkit==0.12.0
-toolz==0.12.1
-torch==2.2.0
-torchdata==0.6.1
-torchdiffeq==0.2.3
-TorchFix==0.4.0
-torchtext==0.15.2
-torchvision==0.17.0
-tqdm==4.66.2
-transformers==4.44.2
-trl==0.9.6
-typer==0.12.5
-types-certifi==2021.10.8.3
-types-requests==2.31.0.20240125
-types-setuptools==69.0.0.20240125
-types-toml==0.10.8.7
-typing==3.7.4.3
-typing-inspect==0.8.0
-typing_extensions==4.9.0
-tyro==0.5.18
-tzdata==2023.3
-unique-names-generator==1.0.2
-urllib3==2.2.2
-uvicorn==0.22.0
-vector_quantize_pytorch==1.14.1
-virtualenv==20.23.0
-voyager==2.0.2
-wandb==0.16.2
-watchfiles==0.21.0
-wavedrom==2.0.3.post3
-wcwidth==0.2.6
-websocket-client==1.7.0
-websockets==12.0
-Werkzeug==3.0.1
-wonderwords==2.2.0
-xxhash==3.2.0
-yarl==1.8.2
-zetascale==2.2.7
-zipp==3.15.0

setup.py

@@ -16,9 +16,7 @@ def parse_requirements():
     with open("./requirements.txt", encoding="utf-8") as requirements_file:
         lines = [r.strip() for r in requirements_file.readlines()]
         for line in lines:
-            is_extras = (
+            is_extras = "deepspeed" in line or "mamba-ssm" in line
-                "deepspeed" in line or "mamba-ssm" in line or "lion-pytorch" in line
-            )
             if line.startswith("--extra-index-url"):
                 # Handle custom index URLs
                 _, url = line.split()
@@ -135,15 +133,15 @@ setup(
         "mlflow": [
             "mlflow",
         ],
-        "lion-pytorch": [
-            "lion-pytorch==0.1.2",
-        ],
         "galore": [
             "galore_torch",
         ],
+        "apollo": [
+            "apollo-torch",
+        ],
         "optimizers": [
             "galore_torch",
-            "lion-pytorch==0.1.2",
+            "apollo-torch",
             "lomo-optim==0.1.1",
             "torch-optimi==0.2.1",
         ],

src/axolotl/cli/inference.py

@@ -56,7 +56,7 @@ def do_inference(
         cfg: Dictionary mapping `axolotl` config keys to values.
         cli_args: Inference-specific CLI arguments.
     """
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg, inference=True)
+    model, tokenizer, _ = load_model_and_tokenizer(cfg=cfg, inference=True)
     prompter = cli_args.prompter
 
     prompter_module = None
@@ -151,7 +151,7 @@ def do_inference_gradio(
     """
     import gradio as gr
 
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg, inference=True)
+    model, tokenizer, _ = load_model_and_tokenizer(cfg=cfg, inference=True)
     prompter = cli_args.prompter
 
     prompter_module = None

src/axolotl/cli/merge_lora.py

@@ -27,7 +27,7 @@ def do_merge_lora(*, cfg: DictDefault) -> None:
     """
     print_axolotl_text_art()
 
-    model, tokenizer = load_model_and_tokenizer(cfg=cfg)
+    model, tokenizer, processor = load_model_and_tokenizer(cfg=cfg)
     safe_serialization = cfg.save_safetensors is True
 
     LOG.info("Running merge of LoRA with base model...")
@@ -44,6 +44,9 @@ def do_merge_lora(*, cfg: DictDefault) -> None:
         )
         tokenizer.save_pretrained(str(Path(cfg.output_dir) / "merged"))
 
+        if processor:
+            processor.save_pretrained(str(Path(cfg.output_dir) / "merged"))
+
 
 def do_cli(config: Union[Path, str] = Path("examples/"), **kwargs) -> None:
     """

src/axolotl/cli/train.py

@@ -17,6 +17,7 @@ from axolotl.cli.config import load_cfg
 from axolotl.common.datasets import load_datasets, load_preference_datasets
 from axolotl.integrations.base import PluginManager
 from axolotl.train import train
+from axolotl.utils import set_pytorch_cuda_alloc_conf
 from axolotl.utils.config import normalize_config, resolve_dtype
 from axolotl.utils.dict import DictDefault
 
@@ -33,6 +34,9 @@ def do_train(cfg: DictDefault, cli_args: TrainerCliArgs):
         cfg: Dictionary mapping `axolotl` config keys to values.
         cli_args: Training-specific CLI arguments.
     """
+    # Enable expandable segments for cuda allocation to improve VRAM usage
+    set_pytorch_cuda_alloc_conf()
+
     print_axolotl_text_art()
     check_accelerate_default_config()
     if int(os.getenv("LOCAL_RANK", "0")) == 0:

src/axolotl/cli/utils.py

@@ -13,11 +13,16 @@ from typing import Any, Callable, Type, Union, get_args, get_origin
 import click
 import requests
 from pydantic import BaseModel
-from transformers import PreTrainedModel, PreTrainedTokenizer, PreTrainedTokenizerFast
+from transformers import (
+    PreTrainedModel,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+    ProcessorMixin,
+)
 
 from axolotl.logging_config import configure_logging
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model, load_tokenizer
+from axolotl.utils.models import load_model, load_processor, load_tokenizer
 
 configure_logging()
 LOG = logging.getLogger(__name__)
@@ -295,9 +300,13 @@ def load_model_and_tokenizer(
     *,
     cfg: DictDefault,
     inference: bool = False,
-) -> tuple[PreTrainedModel, PreTrainedTokenizer | PreTrainedTokenizerFast | Any]:
+) -> tuple[
+    PreTrainedModel,
+    PreTrainedTokenizer | PreTrainedTokenizerFast | Any,
+    ProcessorMixin | None,
+]:
     """
-    Helper function for loading a model and tokenizer specified in the given `axolotl`
+    Helper function for loading a model, tokenizer, and processor specified in the given `axolotl`
     config.
 
     Args:
@@ -305,7 +314,7 @@ def load_model_and_tokenizer(
         inference: Boolean denoting inference mode.
 
     Returns:
-        `transformers` model and tokenizer.
+        Tuple of (PreTrainedModel, PreTrainedTokenizer, ProcessorMixin).
     """
     LOG.info(f"loading tokenizer... {cfg.tokenizer_config or cfg.base_model_config}")
     tokenizer = load_tokenizer(cfg)
@@ -313,4 +322,9 @@ def load_model_and_tokenizer(
     LOG.info("loading model...")
     model, _ = load_model(cfg, tokenizer, inference=inference)
 
-    return model, tokenizer
+    processor = None
+    if cfg.is_multimodal:
+        LOG.info("loading processor...")
+        processor = load_processor(cfg, tokenizer)
+
+    return model, tokenizer, processor

src/axolotl/core/trainer_builder.py

@@ -60,6 +60,7 @@ from axolotl.core.training_args import (
 from axolotl.integrations.base import PluginManager
 from axolotl.monkeypatch.multipack import SUPPORTED_MULTIPACK_MODEL_TYPES
 from axolotl.monkeypatch.relora import ReLoRACallback
+from axolotl.processing_strategies import get_processing_strategy
 from axolotl.utils import is_comet_available, is_mlflow_available
 from axolotl.utils.callbacks import (
     EvalFirstStepCallback,
@@ -747,6 +748,12 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
                 self.cfg.accelerator_config
             )
 
+        if self.cfg.image_size:
+            training_arguments_kwargs["image_size"] = self.cfg.image_size
+        if self.cfg.image_resize_algorithm:
+            training_arguments_kwargs["image_resize_algorithm"] = (
+                self.cfg.image_resize_algorithm
+            )
         if self.cfg.kd_ce_alpha is not None:
             training_arguments_kwargs["kd_ce_alpha"] = self.cfg.kd_ce_alpha
         if self.cfg.kd_alpha is not None:
@@ -890,8 +897,13 @@ class HFCausalTrainerBuilder(TrainerBuilderBase):
         else:
             if self.cfg.processor_type and self.processor:
                 collator = MultiModalChatDataCollator
-                kwargs["processor"] = self.processor
+                kwargs["processing_strategy"] = get_processing_strategy(
-                kwargs["chat_template"] = training_args.chat_template
+                    self.processor,
+                    training_args.chat_template,
+                    self.cfg.chat_template,
+                    image_size=training_args.image_size,
+                    image_resize_algorithm=training_args.image_resize_algorithm,
+                )
             elif self.cfg.batch_flattening:
                 collator = DataCollatorWithFlattening
                 collator_args.pop(0)

src/axolotl/core/trainers/trl.py

@@ -1,5 +1,7 @@
 """Module for TRL PPO trainer"""
 
+from typing import Literal, Union
+
 import torch
 from tqdm import tqdm
 from trl import (
@@ -79,6 +81,78 @@ class AxolotlORPOTrainer(SchedulerMixin, ORPOTrainer):
 
     tag_names = ["axolotl", "orpo"]
 
+    def get_batch_loss_metrics(
+        self,
+        model,
+        batch: dict[str, Union[list, torch.LongTensor]],
+        train_eval: Literal["train", "eval"] = "train",
+    ):
+        """Compute the ORPO loss and other metrics for the given batch of inputs for train or test."""
+
+        # TODO remove once https://github.com/huggingface/trl/pull/3069 is included in a trl release
+
+        metrics = {}
+
+        forward_output = self.concatenated_forward(model, batch)
+        (
+            policy_chosen_logps,
+            policy_rejected_logps,
+            policy_chosen_logits,
+            policy_rejected_logits,
+            policy_nll_loss,
+        ) = forward_output[:5]
+        if self.aux_loss_enabled:
+            aux_loss = forward_output[5]
+
+        losses, chosen_rewards, rejected_rewards, log_odds_ratio, log_odds_chosen = (
+            self.odds_ratio_loss(policy_chosen_logps, policy_rejected_logps)
+        )
+        # full ORPO loss
+        loss = policy_nll_loss - losses.mean()
+
+        reward_accuracies = (chosen_rewards > rejected_rewards).float()
+
+        prefix = "eval_" if train_eval == "eval" else ""
+        metrics[f"{prefix}rewards/chosen"] = self.accelerator.gather_for_metrics(
+            chosen_rewards
+        ).mean()
+        metrics[f"{prefix}rewards/rejected"] = self.accelerator.gather_for_metrics(
+            rejected_rewards
+        ).mean()
+        metrics[f"{prefix}rewards/accuracies"] = self.accelerator.gather_for_metrics(
+            reward_accuracies
+        ).mean()
+        metrics[f"{prefix}rewards/margins"] = self.accelerator.gather_for_metrics(
+            chosen_rewards - rejected_rewards
+        ).mean()
+        metrics[f"{prefix}logps/rejected"] = (
+            self.accelerator.gather_for_metrics(policy_rejected_logps).detach().mean()
+        )
+        metrics[f"{prefix}logps/chosen"] = (
+            self.accelerator.gather_for_metrics(policy_chosen_logps).detach().mean()
+        )
+        metrics[f"{prefix}logits/rejected"] = self.accelerator.gather_for_metrics(
+            policy_rejected_logits.detach().mean()
+        ).mean()
+        metrics[f"{prefix}logits/chosen"] = self.accelerator.gather_for_metrics(
+            policy_chosen_logits.detach().mean()
+        ).mean()
+        metrics[f"{prefix}nll_loss"] = (
+            self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean()
+        )
+        metrics[f"{prefix}log_odds_ratio"] = (
+            self.accelerator.gather_for_metrics(log_odds_ratio).detach().mean()
+        )
+        metrics[f"{prefix}log_odds_chosen"] = (
+            self.accelerator.gather_for_metrics(log_odds_chosen).detach().mean()
+        )
+        for k, v in metrics.items():
+            metrics[k] = v.item()
+        if self.aux_loss_enabled:
+            loss += self.aux_loss_coef * aux_loss
+
+        return loss, metrics
+
 
 class AxolotlKTOTrainer(SchedulerMixin, KTOTrainer):
     """
@@ -95,6 +169,80 @@ class AxolotlCPOTrainer(SchedulerMixin, CPOTrainer):
 
     tag_names = ["axolotl", "cpo"]
 
+    def get_batch_loss_metrics(
+        self,
+        model,
+        batch: dict[str, Union[list, torch.LongTensor]],
+        train_eval: Literal["train", "eval"] = "train",
+    ):
+        """Compute the CPO loss and other metrics for the given batch of inputs for train or test."""
+        metrics = {}
+
+        forward_output = self.concatenated_forward(model, batch)
+        (
+            policy_chosen_logps,
+            policy_rejected_logps,
+            policy_chosen_logits,
+            policy_rejected_logits,
+            policy_nll_loss,
+        ) = forward_output[:5]
+        if self.aux_loss_enabled:
+            aux_loss = forward_output[5]
+
+        losses, chosen_rewards, rejected_rewards = self.cpo_loss(
+            policy_chosen_logps,
+            policy_rejected_logps,
+        )
+
+        loss = losses.mean() + self.cpo_alpha * policy_nll_loss
+        reward_accuracies = (chosen_rewards > rejected_rewards).float()
+
+        prefix = "eval_" if train_eval == "eval" else ""
+        metrics[f"{prefix}rewards/chosen"] = (
+            self.accelerator.gather_for_metrics(chosen_rewards).mean().item()
+        )
+        metrics[f"{prefix}rewards/rejected"] = (
+            self.accelerator.gather_for_metrics(rejected_rewards).mean().item()
+        )
+        metrics[f"{prefix}rewards/accuracies"] = (
+            self.accelerator.gather_for_metrics(reward_accuracies).mean().item()
+        )
+        metrics[f"{prefix}rewards/margins"] = (
+            self.accelerator.gather_for_metrics(chosen_rewards - rejected_rewards)
+            .mean()
+            .item()
+        )
+        metrics[f"{prefix}logps/rejected"] = (
+            self.accelerator.gather_for_metrics(policy_rejected_logps)
+            .detach()
+            .mean()
+            .item()
+        )
+        metrics[f"{prefix}logps/chosen"] = (
+            self.accelerator.gather_for_metrics(policy_chosen_logps)
+            .detach()
+            .mean()
+            .item()
+        )
+        metrics[f"{prefix}logits/rejected"] = (
+            self.accelerator.gather_for_metrics(policy_rejected_logits.detach().mean())
+            .mean()
+            .item()
+        )
+        metrics[f"{prefix}logits/chosen"] = (
+            self.accelerator.gather_for_metrics(policy_chosen_logits.detach().mean())
+            .mean()
+            .item()
+        )
+        metrics[f"{prefix}nll_loss"] = (
+            self.accelerator.gather_for_metrics(policy_nll_loss).detach().mean().item()
+        )
+
+        if self.aux_loss_enabled:
+            loss += self.aux_loss_coef * aux_loss
+
+        return loss, metrics
+
 
 class AxolotlRewardTrainer(SchedulerMixin, RewardTrainer):
     """

src/axolotl/core/training_args.py

@@ -5,6 +5,7 @@ extra axolotl specific training args
 from dataclasses import dataclass, field
 from typing import Optional
 
+from PIL.Image import Resampling
 from transformers import TrainingArguments
 from trl import CPOConfig, KTOConfig, ORPOConfig, PRMConfig, RewardConfig
 
@@ -212,6 +213,20 @@ class AxolotlTrainingMixins:
         metadata={"help": "The number of workers to use in sequence parallelism"},
     )
 
+    # multi-modal section
+
+    image_size: int | tuple[int, int] | None = field(
+        default=None,
+        metadata={"help": "The size of the image to resize to"},
+    )
+
+    image_resize_algorithm: Resampling | None = field(
+        default=None,
+        metadata={"help": "The algorithm to use for image resizing"},
+    )
+
+    # end of multi-modal section
+
 
 @dataclass
 class AxolotlTrainingArguments(AxolotlTrainingMixins, TrainingArguments):

src/axolotl/integrations/cut_cross_entropy/README.md

@@ -1,6 +1,6 @@
 # Cut Cross Entropy
 
-Cut Cross Entropy reduces VRAM usage through optimization on the cross-entropy operation during loss calculation.
+Cut Cross Entropy (CCE) reduces VRAM usage through optimization on the cross-entropy operation during loss calculation.
 
 See https://github.com/apple/ml-cross-entropy
 
@@ -29,6 +29,20 @@ plugins:
 cut_cross_entropy: true
 ```
 
+## Supported Models
+
+- llama
+- phi3
+- gemma
+- gemma2
+- gemma3
+- gemma3_text
+- mistral
+- mistral3
+- qwen2
+- cohere
+- cohere2
+
 ## Citation
 
 ```bib

src/axolotl/integrations/cut_cross_entropy/__init__.py

@@ -72,7 +72,9 @@ class CutCrossEntropyPlugin(BasePlugin):
         if cfg.cut_cross_entropy:
             self._check_requirements()
 
-            from cut_cross_entropy.transformers import cce_patch
+            from axolotl.integrations.cut_cross_entropy.monkeypatch.patch import (
+                cce_patch,
+            )
 
             with zero_only():
                 LOG.info(

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

@@ -0,0 +1,201 @@
+"""Cohere and Cohere2 CCE patch."""
+
+# This patch is based off transformers 4.50.0.
+# It patches the forward function for CohereForCausalLM and Cohere2ForCausalLM.
+# It scales the hidden states by the logit scale in advance instead of the logits as the
+# operation is done internally and should be mathematically equivalent.
+
+# 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 transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.cohere.modeling_cohere import (
+    _CONFIG_FOR_DOC,
+    COHERE_INPUTS_DOCSTRING,
+    KwargsForCausalLM,
+)
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    replace_return_docstrings,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+
+_PATCH_OPTS: PatchOptions | None = None
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(COHERE_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+)
+def cce_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Union[Cache, 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **kwargs: Unpack[KwargsForCausalLM],
+) -> Union[Tuple, 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, CohereForCausalLM
+
+    >> model = CohereForCausalLM.from_pretrained("CohereForAI/c4ai-command-r-v01")
+    >> tokenizer = AutoTokenizer.from_pretrained("CohereForAI/c4ai-command-r-v01")
+
+    >> 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
+    )
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = 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,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    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
+        # scale weight by logit_scale in-place of logits
+        loss = apply_lce(
+            hidden_states[:, slice_indices, :],
+            self.lm_head.weight * self.logit_scale,
+            labels,
+            _PATCH_OPTS,
+            **kwargs,
+        )
+    else:
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        logits = logits * self.logit_scale  # main diff from Llama
+
+        if labels is not None:
+            loss = self.loss_function(
+                logits=logits,
+                labels=labels,
+                vocab_size=self.config.vocab_size,
+                **kwargs,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def patch_cohere(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.cohere import modeling_cohere
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_cohere.CohereForCausalLM
+        ), f"Expected a CohereForCausalLM model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward, maybe_model)
+        return maybe_model
+
+    modeling_cohere.CohereForCausalLM.forward = cce_forward
+    return None
+
+
+def patch_cohere2(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.cohere2 import modeling_cohere2
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_cohere2.Cohere2ForCausalLM
+        ), f"Expected a Cohere2ForCausalLM model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward, maybe_model)
+        return maybe_model
+
+    modeling_cohere2.Cohere2ForCausalLM.forward = cce_forward
+    return None

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

@@ -0,0 +1,175 @@
+"""Gemma CCE patch"""
+
+# This patch is based off transformers 4.50.0.
+
+# 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 transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.gemma.modeling_gemma import (
+    _CONFIG_FOR_DOC,
+    GEMMA_INPUTS_DOCSTRING,
+    KwargsForCausalLM,
+)
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    replace_return_docstrings,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+
+_PATCH_OPTS: PatchOptions | None = None
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(GEMMA_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+)
+def cce_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Union[Cache, 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **kwargs: Unpack[KwargsForCausalLM],
+) -> Union[Tuple, 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, GemmaForCausalLM
+
+    >>> model = GemmaForCausalLM.from_pretrained("google/gemma-7b")
+    >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b")
+
+    >>> prompt = "What is your favorite condiment?"
+    >>> 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]
+    "What is your favorite condiment?"
+    ```"""
+    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
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = 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,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    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,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def patch_gemma(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.gemma import modeling_gemma
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_gemma.GemmaForCausalLM
+        ), f"Expected a GemmaForCausalLM model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward, maybe_model)
+        return maybe_model
+
+    modeling_gemma.GemmaForCausalLM.forward = cce_forward
+    return None

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

@@ -0,0 +1,465 @@
+"""Gemma2 and Gemma3 (text and multimodal) CCE patch."""
+
+# Implementation originally adapted from https://github.com/apple/ml-cross-entropy/pull/29
+# and updated for transformers 4.50.0.
+# This is a modified version of the patch that allows for deferred logits calculation for gemma3 and works
+# with both gemma3 (text and multimodal) models.
+
+# 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 import nn
+from transformers.cache_utils import Cache, HybridCache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.gemma3.modeling_gemma3 import (
+    _CONFIG_FOR_DOC,
+    GEMMA3_INPUTS_DOCSTRING,
+    Gemma3CausalLMOutputWithPast,
+    logger,
+)
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    is_torchdynamo_compiling,
+    replace_return_docstrings,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+
+_PATCH_OPTS: PatchOptions | None = None
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(GEMMA3_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+)
+def cce_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[HybridCache] = 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    defer_logits_calculation: bool = False,
+    **loss_kwargs,
+) -> Union[Tuple, 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).
+
+        defer_logits_calculation (`bool`, *optional*):
+            If `True`, defer logits calculation to the ConditionalGeneration forward. This is used to avoid the
+            memory overhead of calculating logits using regular lm_head forward pass and to use CCE.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, Gemma3ForCausalLM
+
+    >>> model = Gemma3ForCausalLM.from_pretrained("google/gemma-2-9b")
+    >>> tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
+
+    >>> prompt = "What is your favorite condiment?"
+    >>> 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]
+    "What is your favorite condiment?"
+    ```"""
+    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
+    )
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = 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,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        **loss_kwargs,
+    )
+
+    hidden_states = outputs[0]
+    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
+        if self.config.final_logit_softcapping is not None:
+            logger.warning_once(
+                "final_logit_softcapping is not supported for gemma3_text with CCE. Disabling."
+            )
+        loss = apply_lce(
+            hidden_states[:, slice_indices, :],
+            self.lm_head.weight,
+            labels,
+            _PATCH_OPTS,
+            **loss_kwargs,
+        )
+    elif _PATCH_OPTS is not None and defer_logits_calculation:
+        # defer logits calculation to the ConditionalGeneration forward
+        logits = hidden_states[:, slice_indices, :]
+
+        if self.config.final_logit_softcapping is not None:
+            logger.warning_once(
+                "final_logit_softcapping is not supported for gemma3 with CCE. Disabling."
+            )
+    else:
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+        if self.config.final_logit_softcapping is not None:
+            logits = logits / self.config.final_logit_softcapping
+            logits = torch.tanh(logits)
+            logits = logits * self.config.final_logit_softcapping
+
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(GEMMA3_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=Gemma3CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+)
+def cce_forward_multimodal(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    pixel_values: torch.FloatTensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Union[list[torch.FloatTensor], Cache]] = None,
+    token_type_ids: Optional[torch.LongTensor] = None,
+    cache_position: Optional[torch.LongTensor] = 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,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **lm_kwargs,
+) -> Union[Tuple, Gemma3CausalLMOutputWithPast]:
+    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.text_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.text_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 PIL import Image
+    >>> import requests
+    >>> from transformers import AutoProcessor, Gemma3ForConditionalGeneration
+
+    >>> model = Gemma3ForConditionalGeneration.from_pretrained("google/Gemma3-test-224px-hf")
+    >>> processor = AutoProcessor.from_pretrained("google/Gemma3-test-224px-hf")
+
+    >>> prompt = "answer en Where is the cow standing?"
+    >>> url = "https://huggingface.co/gv-hf/Gemma3-test-224px-hf/resolve/main/cow_beach_1.png"
+    >>> image = Image.open(requests.get(url, stream=True).raw)
+
+    >>> inputs = processor(images=image, text=prompt,  return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(**inputs, max_length=30)
+    >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "answer en Where is the cow standing?\nbeach"
+    ```"""
+
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    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
+    )
+
+    is_training = token_type_ids is not None and labels is not None
+
+    # Replace image id woth PAD if the image token if OOV, to avoid index-errors
+    if input_ids is not None and self.config.image_token_index >= self.vocab_size:
+        special_image_mask = input_ids == self.config.image_token_index
+        llm_input_ids = input_ids.clone()
+        llm_input_ids[special_image_mask] = 0
+    else:
+        llm_input_ids = input_ids  # type: ignore
+
+    if inputs_embeds is None:
+        inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+    if cache_position is None:
+        past_seen_tokens = (
+            past_key_values.get_seq_length() if past_key_values is not None else 0  # type: ignore
+        )
+        cache_position = torch.arange(  # type: ignore
+            past_seen_tokens,
+            past_seen_tokens + inputs_embeds.shape[1],
+            device=inputs_embeds.device,
+        )
+
+    # Merge text and images
+    if pixel_values is not None:
+        image_features = self.get_image_features(pixel_values)
+
+        if input_ids is None:
+            special_image_mask = inputs_embeds == self.get_input_embeddings()(
+                torch.tensor(
+                    self.config.image_token_index,
+                    dtype=torch.long,
+                    device=inputs_embeds.device,
+                )
+            )
+        else:
+            special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(
+                -1
+            )
+            special_image_mask = special_image_mask.expand_as(inputs_embeds).to(
+                inputs_embeds.device
+            )
+
+        if (
+            not is_torchdynamo_compiling()
+            and inputs_embeds[special_image_mask].numel() != image_features.numel()
+        ):
+            image_tokens_in_text = (special_image_mask).sum(dim=1).sum(dim=0)[0]
+            raise ValueError(
+                f"Number of images does not match number of special image tokens in the input text. "
+                f"Got {image_tokens_in_text} image tokens in the text but {image_features.shape[0] * image_features.shape[1]} "
+                "tokens from image embeddings."
+            )
+        image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+        inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)  # type: ignore
+
+    # mask out pad-token-ids in labels for BC
+    if labels is not None and self.pad_token_id in labels:
+        logger.warning_once(
+            "`labels` contains `pad_token_id` which will be masked with `config.ignore_index`. "
+            "You have to mask out `pad_token_id` when preparing `labels`, this behavior will be removed in v.4.46.",
+        )
+        labels = torch.where(  # type: ignore
+            input_ids == self.pad_token_id, self.config.ignore_index, labels
+        )
+
+    causal_mask = self._update_causal_mask(  # pylint: disable=protected-access
+        attention_mask,
+        token_type_ids,
+        past_key_values,
+        cache_position,
+        inputs_embeds,
+        is_training,
+    )
+    outputs = self.language_model(
+        attention_mask=causal_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,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        logits_to_keep=logits_to_keep,
+        defer_logits_calculation=True,  # enable deferred logits calculation
+        **lm_kwargs,
+    )
+
+    hidden_states = outputs[0]
+    loss = None
+    logits = 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.language_model.lm_head.weight,
+            labels,
+            _PATCH_OPTS,
+            **lm_kwargs,
+        )
+    else:
+        logits = 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_logits = logits[..., :-1, :]
+            shift_labels = labels[..., 1:]
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -shift_logits.shape[1] :].to(
+                    logits.device
+                )
+                shift_logits = shift_logits[
+                    shift_attention_mask.to(logits.device) != 0
+                ].contiguous()
+                shift_labels = shift_labels[
+                    shift_attention_mask.to(shift_labels.device) != 0
+                ].contiguous()
+            else:
+                shift_logits = shift_logits.contiguous()
+                shift_labels = shift_labels.contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+
+            flat_logits = shift_logits.view(-1, self.config.text_config.vocab_size)
+            flat_labels = shift_labels.view(-1).to(shift_logits.device)
+            loss = loss_fct(flat_logits, flat_labels)
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return Gemma3CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        image_hidden_states=image_features if pixel_values is not None else None,
+    )
+
+
+def patch_gemma2(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.gemma2 import modeling_gemma2
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_gemma2.Gemma2ForCausalLM
+        ), f"Expected a Gemma2ForCausalLM model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward, maybe_model)
+        return maybe_model
+
+    modeling_gemma2.Gemma2ForCausalLM.forward = cce_forward
+    return None
+
+
+def patch_gemma3_text(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.gemma3 import modeling_gemma3
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_gemma3.Gemma3ForCausalLM
+        ), f"Expected a Gemma3ForCausalLM model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward, maybe_model)
+        return maybe_model
+
+    modeling_gemma3.Gemma3ForCausalLM.forward = cce_forward
+    return None
+
+
+def patch_gemma3(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.gemma3 import modeling_gemma3
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_gemma3.Gemma3ForConditionalGeneration
+        ), f"Expected a Gemma3ForConditionalGeneration model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
+
+        # patch the causal model to enable deferred logits calculation
+        maybe_model.language_model.forward = MethodType(
+            cce_forward, maybe_model.language_model
+        )
+        return maybe_model
+
+    modeling_gemma3.Gemma3ForConditionalGeneration.forward = cce_forward_multimodal
+    # patch the causal model to enable deferred logits calculation
+    modeling_gemma3.Gemma3ForCausalLM.forward = cce_forward
+    return None

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

@@ -0,0 +1,392 @@
+"""Mistral and Mistral3 CCE patch."""
+
+# 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 import nn
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.mistral3.modeling_mistral3 import (
+    Mistral3CausalLMOutputWithPast,
+)
+from transformers.models.mistral.modeling_mistral import (
+    _CONFIG_FOR_DOC,
+    MISTRAL_INPUTS_DOCSTRING,
+    KwargsForCausalLM,
+)
+from transformers.processing_utils import Unpack
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    is_torchdynamo_compiling,
+    replace_return_docstrings,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+
+_PATCH_OPTS: PatchOptions | None = None
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(MISTRAL_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
+)
+def cce_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: Optional[torch.Tensor] | None = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_values: Optional[Union[Cache, 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    defer_logits_calculation: bool = False,
+    **kwargs: Unpack[KwargsForCausalLM],
+) -> Union[Tuple, 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).
+
+        defer_logits_calculation (`bool`, *optional*):
+            If `True`, defer logits calculation to the ConditionalGeneration forward. This is used to avoid the
+            memory overhead of calculating logits using regular lm_head forward pass and to use CCE.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, MistralForCausalLM
+
+    >>> model = MistralForCausalLM.from_pretrained("meta-mistral/Mistral-2-7b-hf")
+    >>> tokenizer = AutoTokenizer.from_pretrained("meta-mistral/Mistral-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
+    )
+    return_dict = (
+        return_dict if return_dict is not None else self.config.use_return_dict
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = 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,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        **kwargs,
+    )
+
+    hidden_states = outputs[0]
+    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,
+        )
+    elif _PATCH_OPTS is not None and defer_logits_calculation:
+        # defer logits calculation to the ConditionalGeneration forward
+        logits = hidden_states[:, slice_indices, :]
+    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,
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def cce_forward_multimodal(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    pixel_values: torch.FloatTensor | None = 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,
+    vision_feature_layer: Optional[Union[int, list[int]]] = 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    image_sizes: torch.Tensor | None = None,
+    **lm_kwargs,
+) -> Union[Tuple, Mistral3CausalLMOutputWithPast]:
+    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 PIL import Image
+    >>> import requests
+    >>> from transformers import AutoProcessor, Mistral3ForConditionalGeneration
+
+    >>> model = Mistral3ForConditionalGeneration.from_pretrained("mistralai/Mistral-Small-3.1-24B-Instruct-2503")
+    >>> processor = AutoProcessor.from_pretrained("mistralai/Mistral-Small-3.1-24B-Instruct-2503")
+
+    >>> prompt = "<s>[INST][IMG]What is the image?[/INST]"
+    >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+    >>> image = Image.open(requests.get(url, stream=True).raw)
+
+    >>> inputs = processor(images=image, text=prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(**inputs, max_new_tokens=15)
+    >>> processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    "What is the image?The image depicts two cats lying on a pink blanket."
+    ```"""
+
+    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
+    )
+    vision_feature_layer = (
+        vision_feature_layer
+        if vision_feature_layer is not None
+        else self.config.vision_feature_layer
+    )
+
+    if (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    if pixel_values is not None and inputs_embeds is not None:
+        raise ValueError(
+            "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
+        )
+
+    if inputs_embeds is None:
+        inputs_embeds = self.get_input_embeddings()(input_ids)
+
+    if pixel_values is not None:
+        image_features = self.get_image_features(
+            pixel_values=pixel_values,
+            vision_feature_layer=vision_feature_layer,
+            image_sizes=image_sizes,
+        )
+
+        special_image_mask = (input_ids == self.config.image_token_index).unsqueeze(-1)
+        special_image_mask = special_image_mask.expand_as(inputs_embeds).to(
+            inputs_embeds.device
+        )
+        if (
+            not is_torchdynamo_compiling()
+            and inputs_embeds[special_image_mask].numel() != image_features.numel()
+        ):
+            n_image_tokens = (input_ids == self.config.image_token_index).sum()
+            n_image_features = image_features.shape[0] * image_features.shape[1]
+            raise ValueError(
+                f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
+            )
+        image_features = image_features.to(inputs_embeds.device, inputs_embeds.dtype)
+        inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)  # type: ignore
+
+    outputs = self.language_model(
+        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,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        logits_to_keep=logits_to_keep,
+        defer_logits_calculation=True,  # enable deferred logits calculation
+        **lm_kwargs,
+    )
+
+    hidden_states = outputs[0]
+    loss = None
+    logits = 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.language_model.lm_head.weight,
+            labels,
+            _PATCH_OPTS,
+            **lm_kwargs,
+        )
+    else:
+        logits = hidden_states
+        if labels is not None:
+            # Shift so that tokens < n predict n
+            if attention_mask is not None:
+                # we use the input attention mask to shift the logits and labels, because it is 2D.
+                # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+                shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(
+                    logits.device
+                )
+                shift_logits = logits[..., :-1, :][
+                    shift_attention_mask.to(logits.device) != 0
+                ].contiguous()
+                shift_labels = labels[..., 1:][
+                    shift_attention_mask.to(labels.device) != 0
+                ].contiguous()
+            else:
+                shift_logits = logits[..., :-1, :].contiguous()
+                shift_labels = labels[..., 1:].contiguous()
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1).to(shift_logits.device),
+            )
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return Mistral3CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+        image_hidden_states=image_features if pixel_values is not None else None,
+    )
+
+
+def patch_mistral(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.mistral import modeling_mistral
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_mistral.MistralForCausalLM
+        ), f"Expected a MistralForCausalLM model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward, maybe_model)
+        return maybe_model
+
+    modeling_mistral.MistralForCausalLM.forward = cce_forward
+    return None
+
+
+def patch_mistral3(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.mistral import modeling_mistral
+    from transformers.models.mistral3 import modeling_mistral3
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_mistral3.Mistral3ForConditionalGeneration
+        ), f"Expected a Mistral3ForConditionalGeneration model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
+
+        # patch the causal model to enable deferred logits calculation
+        maybe_model.language_model.forward = MethodType(
+            cce_forward, maybe_model.language_model
+        )
+        return maybe_model
+
+    modeling_mistral3.Mistral3ForConditionalGeneration.forward = cce_forward_multimodal
+    # patch the causal model to enable deferred logits calculation
+    modeling_mistral.MistralForCausalLM.forward = cce_forward
+    return None

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

@@ -0,0 +1,379 @@
+"""Mllama CCE patch."""
+
+# 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 transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.mllama.modeling_mllama import (
+    MLLAMA_INPUTS_DOCSTRING,
+    _prepare_cross_attention_mask,
+)
+from transformers.utils import (
+    add_start_docstrings_to_model_forward,
+    replace_return_docstrings,
+)
+from transformers.utils.deprecation import deprecate_kwarg
+
+_PATCH_OPTS: PatchOptions | None = None
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(MLLAMA_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class="MllamaTextConfig"
+)
+def cce_forward(
+    self,
+    input_ids: torch.LongTensor | None = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    cross_attention_states: Optional[torch.LongTensor] = None,
+    cross_attention_mask: Optional[torch.LongTensor] = None,
+    full_text_row_masked_out_mask: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    past_key_values: Optional[Union[Cache, 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    defer_logits_calculation: bool = False,
+    **loss_kwargs,
+) -> Union[Tuple, 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).
+
+        defer_logits_calculation (`bool`, *optional*):
+            If `True`, defer logits calculation to the ConditionalGeneration forward. This is used to avoid the
+            memory overhead of calculating logits using regular lm_head forward pass and to use CCE.
+
+    Returns:
+
+    Example:
+
+    ```python
+    >>> from transformers import AutoTokenizer, MllamaForCausalLM
+
+    >>> model = MllamaForCausalLM.from_pretrained("Llama-3.2-11B-Vision")
+    >>> tokenizer = AutoTokenizer.from_pretrained("Llama-3.2-11B-Vision")
+
+    >>> prompt = "If I had to write a haiku, it would be:"
+    >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+    >>> # Generate
+    >>> generate_ids = model.generate(inputs.input_ids, max_length=40, do_sample=True, temperature=0.6)
+    >>> result = tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    >>> print(result)
+    If I had to write a haiku, it would be: "Snowflakes gently fall" - simple, yet peaceful.
+    I love the idea of snowflakes gently falling, each one
+    ```
+    """
+    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
+    )
+
+    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+    outputs = self.model(
+        input_ids=input_ids,
+        cross_attention_states=cross_attention_states,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        cross_attention_mask=cross_attention_mask,
+        full_text_row_masked_out_mask=full_text_row_masked_out_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]
+    loss = None
+    logits = None
+
+    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,
+        )
+    elif _PATCH_OPTS is not None and defer_logits_calculation:
+        # defer logits calculation to the ConditionalGeneration forward
+        logits = hidden_states[:, slice_indices, :]
+    else:
+        logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits, labels, self.vocab_size, **loss_kwargs)
+
+    if not return_dict:
+        output = (logits,) + outputs[1:]
+        return (loss,) + output if loss is not None else output
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+@deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")
+@add_start_docstrings_to_model_forward(MLLAMA_INPUTS_DOCSTRING)
+@replace_return_docstrings(
+    output_type=CausalLMOutputWithPast, config_class="MllamaConfig"
+)
+def cce_forward_multimodal(
+    self,
+    input_ids: Optional[torch.LongTensor] = None,
+    pixel_values: Optional[torch.FloatTensor] = None,
+    aspect_ratio_mask: Optional[torch.Tensor] = None,
+    aspect_ratio_ids: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+    cross_attention_mask: Optional[torch.Tensor] = None,
+    cross_attention_states: 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,
+    cache_position: Optional[torch.LongTensor] = None,
+    logits_to_keep: Union[int, torch.Tensor] = 0,
+    **loss_kwargs,
+) -> Union[Tuple, 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 PIL import Image
+    >>> import requests
+    >>> from transformers import AutoProcessor, MllamaForConditionalGeneration
+
+    >>> checkpoint = "meta-llama/Llama-3.2-11B-Vision"
+    >>> model = MllamaForConditionalGeneration.from_pretrained(checkpoint)
+    >>> processor = AutoProcessor.from_pretrained(checkpoint)
+
+    >>> prompt = "<|image|>If I had to write a haiku for this one"
+    >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+    >>> image = Image.open(requests.get(url, stream=True).raw)
+
+    >>> inputs = processor(text=prompt, images=image, return_tensors="pt")
+
+    >>> # Generate
+    >>> output = model.generate(**inputs, max_new_tokens=15)
+
+    >>> prompt_len = inputs.input_ids.shape[-1]
+    >>> generated_ids = output[:, prompt_len:]
+    >>> generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+    >>> print(generated_text)
+    [', it would be:.\\nA stop sign in Chinatown.\\n']
+    ```
+    """
+    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 (input_ids is None) ^ (inputs_embeds is not None):
+        raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+    if pixel_values is not None and inputs_embeds is not None:
+        raise ValueError(
+            "You cannot specify both pixel_values and inputs_embeds at the same time, and must specify either one"
+        )
+
+    if pixel_values is not None and cross_attention_states is not None:
+        raise ValueError(
+            "`pixel_values` and `cross_attention_states` cannot be provided simultaneously"
+        )
+
+    if pixel_values is not None:
+        if aspect_ratio_ids is None:
+            raise ValueError(
+                "`aspect_ratio_ids` must be provided if `pixel_values` is provided"
+            )
+        # get vision tokens from vision model
+        vision_outputs = self.vision_model(
+            pixel_values=pixel_values,
+            aspect_ratio_ids=aspect_ratio_ids,
+            aspect_ratio_mask=aspect_ratio_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+        )
+        cross_attention_states = vision_outputs[0]
+        cross_attention_states = self.multi_modal_projector(
+            cross_attention_states
+        ).reshape(
+            -1, cross_attention_states.shape[-2], self.hidden_size  # type: ignore
+        )
+
+    if cross_attention_mask is not None:
+        cross_attention_mask, full_text_row_masked_out_mask = (
+            _prepare_cross_attention_mask(
+                cross_attention_mask,
+                num_vision_tokens=self.vision_model.num_patches,
+                dtype=self.dtype,
+            )
+        )
+    else:
+        full_text_row_masked_out_mask = None
+
+    if cross_attention_mask is not None and cache_position is not None:
+        cross_attention_mask = cross_attention_mask[:, :, cache_position]
+        full_text_row_masked_out_mask = full_text_row_masked_out_mask[
+            :, :, cache_position
+        ]
+
+    outputs = self.language_model(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        position_ids=position_ids,
+        cross_attention_states=cross_attention_states,
+        cross_attention_mask=cross_attention_mask,
+        full_text_row_masked_out_mask=full_text_row_masked_out_mask,
+        past_key_values=past_key_values,
+        use_cache=use_cache,
+        inputs_embeds=inputs_embeds,
+        output_hidden_states=output_hidden_states,
+        output_attentions=output_attentions,
+        return_dict=return_dict,
+        cache_position=cache_position,
+        logits_to_keep=logits_to_keep,
+        defer_logits_calculation=True,  # enable deferred logits calculation
+        **loss_kwargs,
+    )
+
+    hidden_states = outputs[0]
+    loss = None
+    logits = 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.language_model.lm_head.weight,
+            labels,
+            _PATCH_OPTS,
+            **loss_kwargs,
+        )
+    else:
+        # Temporary fix to calculate the loss in main class, as the model's vocab size may be resized
+        logits = hidden_states
+
+        if labels is not None:
+            loss = self.loss_function(
+                logits, labels, self.config.get_text_config().vocab_size, **loss_kwargs
+            )
+
+    if not return_dict:
+        return (loss,) + outputs if loss is not None else outputs
+
+    return CausalLMOutputWithPast(
+        loss=loss,
+        logits=outputs.logits,
+        past_key_values=outputs.past_key_values,
+        hidden_states=outputs.hidden_states,
+        attentions=outputs.attentions,
+    )
+
+
+def patch_mllama(
+    maybe_model: TransformersModelT | str | transformers.PretrainedConfig,
+    patch_options: PatchOptions,
+) -> TransformersModelT | None:
+
+    global _PATCH_OPTS  # pylint: disable=global-statement
+    from transformers.models.mllama import modeling_mllama
+
+    _PATCH_OPTS = patch_options
+
+    if isinstance(maybe_model, transformers.PreTrainedModel):
+        assert isinstance(
+            maybe_model, modeling_mllama.MllamaForConditionalGeneration
+        ), f"Expected a MllamaForConditionalGeneration model. Got {type(maybe_model)}."
+        maybe_model.forward = MethodType(cce_forward_multimodal, maybe_model)
+
+        # patch the language model
+        maybe_model.language_model.forward = MethodType(
+            cce_forward, maybe_model.language_model
+        )
+        return maybe_model
+
+    modeling_mllama.MllamaForConditionalGeneration.forward = cce_forward_multimodal
+
+    # patch the causal language model
+    modeling_mllama.MllamaForCausalLM.forward = cce_forward
+    return None

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

@@ -0,0 +1,85 @@
+# Copyright (C) 2024 Apple Inc. All Rights Reserved.
+
+"""Cut Cross Entropy patcher"""
+
+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 (
+    patch_cohere,
+    patch_cohere2,
+)
+from axolotl.integrations.cut_cross_entropy.monkeypatch.gemma import patch_gemma
+from axolotl.integrations.cut_cross_entropy.monkeypatch.gemma3 import (
+    patch_gemma2,
+    patch_gemma3,
+    patch_gemma3_text,
+)
+from axolotl.integrations.cut_cross_entropy.monkeypatch.mistral3 import (
+    patch_mistral,
+    patch_mistral3,
+)
+from axolotl.integrations.cut_cross_entropy.monkeypatch.mllama import patch_mllama
+
+CUT_CROSS_ENTROPY_MODEL_MAPPING = {
+    "llama": patch_llama,
+    "mllama": patch_mllama,
+    "phi3": patch_phi3,
+    "gemma": patch_gemma,
+    "gemma2": patch_gemma2,
+    "gemma3": patch_gemma3,
+    "gemma3_text": patch_gemma3_text,
+    "mistral": patch_mistral,
+    "mistral3": patch_mistral3,
+    "qwen2": patch_qwen2,
+    "cohere": patch_cohere,
+    "cohere2": patch_cohere2,
+}
+
+
+def cce_patch(
+    model_type_or_model: str | TransformersModelT | transformers.PretrainedConfig,
+    impl: str | LinearCrossEntropyImpl = LCE_IMPL_DEFAULT,
+    reduction: str = "mean",
+    filter_eps: float | str | None = "auto",
+    accum_e_fp32: bool = False,
+    accum_c_fp32: bool = False,
+    filter_e_grad: bool = True,
+    filter_c_grad: bool = True,
+    train_only: bool = False,
+) -> TransformersModelT | None:
+    if isinstance(impl, LinearCrossEntropyImpl):
+        impl = impl.name.lower()
+
+    if impl not in (v.name.lower() for v in LinearCrossEntropyImpl):
+        raise ValueError(f"Unknown {impl=}")
+
+    if isinstance(model_type_or_model, transformers.PreTrainedModel):
+        model_type = model_type_or_model.config.model_type
+    elif isinstance(model_type_or_model, transformers.PretrainedConfig):
+        model_type = model_type_or_model.model_type
+    else:
+        model_type = model_type_or_model
+
+    patch_options = PatchOptions(
+        impl=impl,
+        reduction=reduction,
+        filter_eps=filter_eps,
+        accum_e_fp32=accum_e_fp32,
+        accum_c_fp32=accum_c_fp32,
+        filter_e_grad=filter_e_grad,
+        filter_c_grad=filter_c_grad,
+        train_only=train_only,
+    )
+
+    if model_type in CUT_CROSS_ENTROPY_MODEL_MAPPING:
+        return CUT_CROSS_ENTROPY_MODEL_MAPPING[model_type](
+            model_type_or_model, patch_options
+        )
+
+    raise RuntimeError(f"Unknown model type {model_type}")

src/axolotl/integrations/liger/__init__.py

@@ -114,3 +114,5 @@ class LigerPlugin(BasePlugin):
                 modeling_mod.CrossEntropyLoss = LigerCrossEntropyLoss
             if cfg.liger_fused_linear_cross_entropy:
                 modeling_mod.DeepseekV2ForCausalLM.forward = deepseekv2_lce_forward
+        elif cfg.model_config_type in ["gemma3_text", "deepseek_v3"]:
+            raise ValueError(f"Unsupported model config type: {cfg.model_config_type}")

src/axolotl/monkeypatch/attention/ring_attn.py

@@ -27,7 +27,7 @@ def get_ring_attn_group() -> dist.ProcessGroup:
     return RING_ATTN_GROUP
 
 
-def set_ring_attn_group(ring_attn_group: dist.ProcessGroup):
+def set_ring_attn_group(ring_attn_group: dist.ProcessGroup | None):
     """
     Setter for ring attention group on this rank.
 

src/axolotl/monkeypatch/multipack.py

@@ -22,6 +22,9 @@ SUPPORTED_MULTIPACK_MODEL_TYPES = [
     "phi3",
     "gemma",
     "gemma2",
+    "gemma3_text",
+    "cohere",
+    "cohere2",
     "gemmoe",
     "starcoder2",
     "deepseek_v2",

src/axolotl/processing_strategies.py

@@ -0,0 +1,313 @@
+"""Module containing ProcessingStrategy classes and its derivative for different MultiModal Model types"""
+
+import ast
+from copy import deepcopy
+from typing import Optional
+
+from PIL import Image, ImageOps
+from PIL.Image import Resampling
+from torch import Tensor
+from transformers import ProcessorMixin
+from transformers.image_utils import load_image
+
+
+class ProcessingStrategy:
+    """Base Processing Strategy class"""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+    ):
+        self.processor = processor
+        self.chat_template = chat_template
+        self.image_token = None
+        self.image_token_id = None
+
+        self.image_size = image_size
+        self.image_resize_algorithm = (
+            image_resize_algorithm or Image.Resampling.BILINEAR
+        )
+
+        if hasattr(processor, "image_token"):
+            self.image_token = processor.image_token
+            self.image_token_id = processor.tokenizer.convert_tokens_to_ids(
+                self.image_token
+            )
+
+    def __call__(self, examples: list[dict]) -> list[dict]:
+        """
+        Preprocess conversation examples to ensure consistent format.
+        Converts different conversation formats to OpenAI format with 'messages'.
+        Supports two formats:
+        1. OpenAI format with 'messages'
+        2. Legacy format with 'conversations'
+
+        Args:
+            examples: list of conversation dictionaries
+
+        Returns:
+            list of dicts in OpenAI format with 'messages' key
+
+        Raises:
+            ValueError: If the conversation format is not supported
+        """
+        role_mapping = {
+            "human": "user",
+            "gpt": "assistant",
+        }
+
+        def normalize_role(role: str) -> str:
+            """Normalize role names to OpenAI format. Default to original role if not found."""
+            return role_mapping.get(role, role)
+
+        def convert_legacy_format(example: dict) -> dict:
+            """Convert legacy 'conversations' format to OpenAI 'messages' format."""
+            messages = [
+                {"role": normalize_role(convo["from"]), "content": convo["value"]}
+                for convo in example["conversations"]
+            ]
+
+            # Create new dict without 'conversations' key
+            result = deepcopy(example)
+            result.pop("conversations")
+            result["messages"] = messages
+            return result
+
+        def convert_multiple_choice_to_multimedia_messages(
+            messages: dict,
+        ) -> list[dict]:
+
+            def construct_prompt(sample):
+                question = sample["question"]
+                options = sample["options"]
+                if isinstance(options, str):
+                    options = ast.literal_eval(options)
+
+                example = ""
+                start_chr = "A"
+                prediction_range = []
+                index2ans = {}
+                for option in options:
+                    prediction_range.append(start_chr)
+                    example += f"({start_chr}) {option}\n"
+                    index2ans[start_chr] = option
+                    start_chr = chr(ord(start_chr) + 1)
+
+                empty_prompt_sample_structure = "{}\n\n{}\n\nAnswer with the option's letter from the given choices directly."
+                empty_prompt = empty_prompt_sample_structure.format(question, example)
+
+                return empty_prompt
+
+            new_messages = []
+
+            user_content = construct_prompt(messages)
+            assistant_response = messages["answer"]
+
+            new_messages.append(
+                {"role": "user", "content": [{"type": "text", "text": user_content}]}
+            )
+
+            new_messages.append(
+                {
+                    "role": "assistant",
+                    "content": [{"type": "text", "text": assistant_response}],
+                }
+            )
+
+            return new_messages
+
+        def convert_messages_to_multimedia_messages(messages: list[dict]) -> list[dict]:
+            """Convert regular messages format to Messages format with content type"""
+
+            new_messages = []
+            for message in messages:
+                if isinstance(message["content"], str):
+                    new_messages.append(
+                        {
+                            "role": message["role"],
+                            "content": [
+                                {
+                                    "type": "text",
+                                    "text": message["content"],
+                                }
+                            ],
+                        }
+                    )
+                elif isinstance(message["content"], list):
+                    content = message["content"]
+
+                    new_messages.append(
+                        {
+                            "role": message["role"],
+                            "content": content,
+                        }
+                    )
+
+            return new_messages
+
+        processed_examples = []
+        for example in examples:
+            if not (
+                "messages" in example
+                or "conversations" in example
+                or "question" in example
+            ):
+                raise ValueError(
+                    "Only `messages`, `conversations`, and `question` message keys are currently supported."
+                )
+
+            processed_example = None
+            if "messages" in example:  # OpenAI format
+                processed_example = example
+                # convert regular messages format to Messages format with content type
+                # for compatibility with apply_chat_template
+                processed_example["messages"] = convert_messages_to_multimedia_messages(
+                    processed_example["messages"]
+                )
+            elif "question" in example:  # Multiple choice format
+                processed_example = {}
+                processed_example["messages"] = (
+                    convert_multiple_choice_to_multimedia_messages(example)
+                )
+            else:  # Legacy format
+                processed_example = convert_legacy_format(example)
+                processed_example["messages"] = convert_messages_to_multimedia_messages(
+                    processed_example["messages"]
+                )
+
+            # find the image key if it exists
+
+            image_keys = []
+            for key in example.keys():
+                if "image" in key:
+                    image_keys.append(key)
+
+            for im_key in image_keys:
+                if example[im_key] is None:
+                    continue
+                if isinstance(example[im_key], list):
+                    if len(example[im_key]) == 0:
+                        continue
+                    image_value = example[im_key][0]
+                else:
+                    image_value = example[im_key]
+
+                image_value = load_image(image_value)
+
+                if self.image_size is not None:
+                    assert hasattr(
+                        image_value, "resize"
+                    ), "Image does not have a resize method"
+
+                    if isinstance(self.image_size, tuple):
+                        image_value = image_value.resize(
+                            self.image_size, self.image_resize_algorithm
+                        )
+                    else:
+                        # Set the padding value; here we use black (0, 0, 0) for RGB images
+                        padding_color = (0, 0, 0)
+
+                        # When image_size is an int (square target), preserve aspect ratio then pad
+                        # This is to prevent aspect ratio distortion when resizing to square
+                        image_value = ImageOps.pad(
+                            image_value,
+                            (self.image_size, self.image_size),
+                            method=self.image_resize_algorithm,
+                            color=padding_color,
+                        )
+
+                processed_example["messages"][0]["content"].append(
+                    {
+                        "type": "image",
+                        "image": image_value,
+                    }
+                )
+
+            processed_examples.append(processed_example)
+
+        return processed_examples
+
+    def process_labels(self, input_ids: Tensor) -> Tensor:
+        labels = input_ids.clone()
+
+        # The labels are the input_ids, and we mask the padding tokens in the loss computation
+        labels[labels == self.processor.tokenizer.pad_token_id] = -100
+
+        # Ignore the image token index in the loss computation (model specific)
+        labels[labels == self.image_token_id] = -100
+
+        return labels
+
+
+class Qwen2VLProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Qwen2-VL"""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+    ):
+        super().__init__(processor, chat_template, image_size, image_resize_algorithm)
+        self.image_token = "<|image_pad|>"  # nosec
+        self.image_token_id = processor.tokenizer.convert_tokens_to_ids(
+            self.image_token
+        )
+
+
+class Gemma3ProcessingStrategy(ProcessingStrategy):
+    """Processing Strategy class for Gemma3"""
+
+    def __init__(
+        self,
+        processor: ProcessorMixin,
+        chat_template: Optional[str] = None,
+        image_size: int | tuple[int, int] | None = None,
+        image_resize_algorithm: Resampling | None = None,
+    ):
+        super().__init__(processor, chat_template, image_size, image_resize_algorithm)
+        self.image_token = processor.tokenizer.special_tokens_map["boi_token"]
+        self.image_token_id = processor.tokenizer.convert_tokens_to_ids(
+            self.image_token
+        )
+
+    def process_labels(self, input_ids):
+        labels = input_ids.clone()
+
+        # Follows https://ai.google.dev/gemma/docs/core/huggingface_vision_finetune_qlora
+        labels[labels == self.processor.tokenizer.pad_token_id] = -100
+        labels[labels == self.image_token_id] = -100
+        labels[labels == 262144] = -100  # corresponds to <image_soft_token>
+
+        return labels
+
+
+def get_processing_strategy(
+    processor: ProcessorMixin,
+    chat_template,
+    chat_template_type,
+    image_size: int | tuple[int, int] | None = None,
+    image_resize_algorithm: Resampling | None = None,
+):
+    if chat_template_type == "qwen2_vl":
+        return Qwen2VLProcessingStrategy(
+            processor, chat_template, image_size, image_resize_algorithm
+        )
+    if chat_template_type == "gemma3":
+        return Gemma3ProcessingStrategy(
+            processor, chat_template, image_size, image_resize_algorithm
+        )
+    if chat_template_type in [
+        "llama3_2_vision",
+        "llava",
+        "mistral_v7_tekken",
+        "pixtral",
+    ]:
+        return ProcessingStrategy(
+            processor, chat_template, image_size, image_resize_algorithm
+        )
+    raise ValueError(f"Unsupported chat template type: {chat_template_type}")

src/axolotl/train.py

@@ -314,6 +314,7 @@ def save_initial_configs(
     tokenizer: PreTrainedTokenizer,
     model: PreTrainedModel,
     peft_config: PeftConfig | None,
+    processor: ProcessorMixin | None,
 ):
     """
     Save initial configurations before training.
@@ -341,6 +342,10 @@ def save_initial_configs(
         LOG.info(f"Pre-saving model config to {cfg.output_dir}...")
         model.config.save_pretrained(str(output_dir))
 
+    if processor:
+        LOG.info(f"Pre-saving processor to {cfg.output_dir}...")
+        processor.save_pretrained(str(output_dir))
+
 
 def setup_model_card(cfg: DictDefault):
     """
@@ -408,6 +413,7 @@ def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) ->
     PeftModel | PreTrainedModel,
     PreTrainedTokenizer,
     PeftConfig | None,
+    ProcessorMixin | None,
 ]:
     """
     Load model, tokenizer, trainer, etc. Helper function to encapsulate the full
@@ -423,6 +429,7 @@ def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) ->
             - Model
             - Tokenizer
             - PEFT config
+            - Processor
     """
     # Load tokenizer, processor and model
     model, tokenizer, peft_config, processor = setup_model_and_tokenizer(cfg)
@@ -453,6 +460,7 @@ def setup_model_and_trainer(cfg: DictDefault, dataset_meta: TrainDatasetMeta) ->
         model,
         tokenizer,
         peft_config,
+        processor,
     )
 
 
@@ -475,6 +483,7 @@ def train(
         model,
         tokenizer,
         peft_config,
+        processor,
     ) = setup_model_and_trainer(cfg, dataset_meta)
 
     # Determine if we need to resume from a checkpoint
@@ -490,7 +499,7 @@ def train(
     )
 
     # Save initial configs
-    save_initial_configs(cfg, tokenizer, model, peft_config)
+    save_initial_configs(cfg, tokenizer, model, peft_config, processor)
 
     # Set up signal handler for graceful termination
     setup_signal_handler(cfg, model, safe_serialization)

src/axolotl/utils/collators/mm_chat.py

@@ -2,15 +2,17 @@
 Collators for multi-modal chat messages and packing
 """
 
-from copy import deepcopy
 from dataclasses import dataclass
 from typing import Any, Optional, Union
 
-from PIL import Image
+import torch
-from transformers import PreTrainedTokenizerBase, ProcessorMixin
+from torch import Tensor
+from transformers import PreTrainedTokenizerBase
 from transformers.data.data_collator import DataCollatorMixin
 from transformers.utils import PaddingStrategy
 
+from axolotl.processing_strategies import ProcessingStrategy
+
 
 @dataclass
 class MultiModalChatDataCollator(DataCollatorMixin):
@@ -19,11 +21,9 @@ class MultiModalChatDataCollator(DataCollatorMixin):
     """
 
     tokenizer: PreTrainedTokenizerBase
-    processor: ProcessorMixin
+    processing_strategy: ProcessingStrategy
-    return_tensors: str = "pt"
-    chat_template: Optional[str] = None
     packing: bool = False
-    max_images: int = -1
+    return_tensors: str = "pt"
     padding: Union[bool, str, PaddingStrategy] = True
     pad_to_multiple_of: Optional[int] = None
 
@@ -31,162 +31,62 @@ class MultiModalChatDataCollator(DataCollatorMixin):
         if self.packing:
             raise ValueError("Packing is currently not supported.")
 
-    def torch_call(
+    def torch_call(self, examples: list[dict]) -> dict[str, Any]:
-        self, examples: list[Union[list[int], Any, dict[str, Any]]]
+        return self.process_rows(examples)
-    ) -> dict[str, Any]:
-        # Handle dict or lists with proper padding and conversion to tensor.
-
-        return self.__class__.process_rows(
-            examples, self.processor, self.chat_template, self.max_images
-        )
-
-    @staticmethod
-    def process_rows(examples, processor, chat_template, max_images, length_only=False):
-        # HINT: use `_torch_collate_batch` to stack and pad tensors
-        # see also DataCollatorWithFlattening and DefaultDataCollator
-
-        # *** This is COPIED from the trl example sft_vlm.py code ***
-        # use this as a starting point
-
-        def _preprocess(examples: list[dict]) -> list[dict]:
-            """
-            Preprocess conversation examples to ensure consistent format.
-
-            Converts different conversation formats to OpenAI format with 'messages'.
-            Supports two formats:
-            1. OpenAI format with 'messages'
-            2. Legacy format with 'conversations'
-
-            Args:
-                examples: list of conversation dictionaries
-
-            Returns:
-                dict in OpenAI format with 'messages' key
-
-            Raises:
-                ValueError: If the conversation format is not supported
-            """
-            role_mapping = {
-                "human": "user",
-                "gpt": "assistant",
-            }
-
-            def normalize_role(role: str) -> str:
-                """Normalize role names to OpenAI format. Default to original role if not found."""
-                return role_mapping.get(role, role)
-
-            def convert_legacy_format(example: dict) -> dict:
-                """Convert legacy 'conversations' format to OpenAI 'messages' format."""
-                messages = [
-                    {
-                        "role": normalize_role(convo["from"]),
-                        "content": convo["value"],
-                    }
-                    for convo in example["conversations"]
-                ]
-
-                # Create new dict without 'conversations' key
-                result = deepcopy(example)
-                result.pop("conversations")
-                return {"messages": messages, **result}
-
-            processed_examples = []
-            for example in examples:
-                # OpenAI format
-                if "messages" in example:
-                    processed_examples.append(example)
-
-                # Legacy format
-                elif "conversations" in example:
-                    processed_examples.append(convert_legacy_format(example))
-
-                else:
-                    raise ValueError(
-                        "Only `messages` and `conversations` message keys are currently supported."
-                    )
-
-            return processed_examples
-
-        def _process_images(examples, max_images):
-            """
-            Process images from examples, ensuring consistency in image presence and applying max_images limit.
-
-            Args:
-                examples: List of dictionaries that may contain 'images' key
-                max_images: Maximum number of images to keep per example (0 means no limit)
-
-            Returns:
-                Either None (if no images) or List[Image objects] (if all examples have images)
-
-            Raises:
-                ValueError: If there's a mix of None and non-None images
-            """
-
-            def get_image(example):
-                if "images" not in example:
-                    return None
-                images = example["images"]
-                if isinstance(images, str):
-                    return Image.open(images)
-                return images
-
-            images = [get_image(example) for example in examples]
-
-            # Count None and non-None images
-            none_count = sum(1 for img in images if img is None)
-
-            # All images are None
-            if none_count == len(images):
-                return None
-
-            # Mix of None and non-None images
-            if none_count > 0:
-                raise ValueError(
-                    "All images should be either None or not None. "
-                    "Please provide images for all examples or None."
-                )
-
-            # Apply max_images limit if specified
-            if max_images > 0:
-                images = [
-                    (
-                        img_batch[:max_images]
-                        if isinstance(img_batch, (list, tuple))
-                        else img_batch
-                    )
-                    for img_batch in images
-                ]
-
-            return images
 
+    def process_rows(
+        self,
+        examples: list[dict],
+    ) -> dict[str, Tensor]:
         # Preprocess the examples
-        examples = _preprocess(examples)
+        examples = self.processing_strategy(examples)
 
-        # Get the texts and images, and apply the chat template
+        # Initialize batch
-        texts = [
+        batch: dict[str, Any] = {}
-            processor.apply_chat_template(
+
-                example["messages"], chat_template=chat_template, tokenize=False
+        # Process each example
+        for example in examples:
+            # Apply chat template to process the example
+            # This method requires transformers>=4.49.0
+            result = self.processing_strategy.processor.apply_chat_template(
+                example["messages"],
+                add_generation_prompt=True,
+                tokenize=True,
+                return_tensors="pt",
+                padding=True,
+                return_dict=True,
+                chat_template=self.processing_strategy.chat_template,
             )
-            for example in examples
-        ]
 
-        images = _process_images(examples, max_images=max_images)
+            # TODO: Check if need handling for len(input_ids) > sequence_len
 
-        # Tokenize the texts and process the images
+            # Add the processed tensors to our batch
-        batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
+            for key in result.keys():
+                if key not in batch:
+                    batch[key] = []
 
-        # The labels are the input_ids, and we mask the padding tokens in the loss computation
+                batch[key].append(result[key].squeeze(0))
-        labels = batch["input_ids"].clone()
+
-        labels[labels == processor.tokenizer.pad_token_id] = -100  #
+        # Pad sequences to the same length
-        # Ignore the image token index in the loss computation (model specific)
+        input_ids = torch.nn.utils.rnn.pad_sequence(
-        image_token_id = processor.tokenizer.convert_tokens_to_ids(
+            batch["input_ids"],
-            processor.image_token
+            batch_first=True,
+            padding_value=self.tokenizer.pad_token_id,
         )
-        labels[labels == image_token_id] = -100
-        batch["labels"] = labels
 
-        if length_only:
+        attention_mask = torch.nn.utils.rnn.pad_sequence(
-            return {
+            batch["attention_mask"], batch_first=True, padding_value=0
-                "length": [len(sample["input_ids"]) for sample in batch["input_ids"]]
+        )
-            }
+
-        return batch
+        # Create the final batch
+        final_batch = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+        }
+
+        # Process the labels
+        final_batch["labels"] = self.processing_strategy.process_labels(
+            final_batch["input_ids"]
+        )
+
+        return final_batch

src/axolotl/utils/config/__init__.py

@@ -13,7 +13,7 @@ from axolotl.integrations.base import PluginManager
 from axolotl.integrations.config import merge_input_args
 from axolotl.utils.bench import log_gpu_memory_usage
 from axolotl.utils.dict import DictDefault
-from axolotl.utils.models import load_model_config
+from axolotl.utils.models import MULTIMODAL_AUTO_MODEL_MAPPING, load_model_config
 from axolotl.utils.schemas.config import (
     AxolotlConfigWCapabilities as AxolotlConfigWCapabilitiesBase,
 )
@@ -158,7 +158,7 @@ def normalize_config(cfg):
 
     cfg.is_multimodal = (
         hasattr(model_config, "model_type")
-        and model_config.model_type in ["llava", "mllama"]
+        and model_config.model_type in MULTIMODAL_AUTO_MODEL_MAPPING
         or any(
             multimodal_name in cfg.base_model.lower()
             for multimodal_name in [
@@ -171,7 +171,6 @@ def normalize_config(cfg):
         cfg.processor_config = (
             cfg.processor_config or cfg.base_model_config or cfg.base_model
         )
-        model_config = model_config.text_config
 
     cfg.model_config_type = model_config.model_type
 

src/axolotl/utils/models.py

@@ -34,12 +34,16 @@ from transformers import (  # noqa: F401
     AutoTokenizer,
     AwqConfig,
     BitsAndBytesConfig,
+    Gemma3ForConditionalGeneration,
     GPTQConfig,
     LlavaForConditionalGeneration,
+    Mistral3ForConditionalGeneration,
     MllamaForConditionalGeneration,
     PreTrainedModel,
     PreTrainedTokenizerBase,
     ProcessorMixin,
+    Qwen2_5_VLForConditionalGeneration,
+    Qwen2VLForConditionalGeneration,
 )
 from transformers.integrations.deepspeed import (
     HfTrainerDeepSpeedConfig,
@@ -69,9 +73,13 @@ from axolotl.utils.model_shard_quant import load_sharded_model, load_sharded_mod
 
 LOG = logging.getLogger(__name__)
 
-MULTIMODEL_AUTO_MODEL_MAPPING = {
+MULTIMODAL_AUTO_MODEL_MAPPING = {
-    "llava": LlavaForConditionalGeneration,
     "mllama": MllamaForConditionalGeneration,
+    "llava": LlavaForConditionalGeneration,
+    "qwen2_vl": Qwen2VLForConditionalGeneration,
+    "qwen2_5_vl": Qwen2_5_VLForConditionalGeneration,
+    "mistral3": Mistral3ForConditionalGeneration,
+    "gemma3": Gemma3ForConditionalGeneration,
 }
 
 
@@ -101,7 +109,21 @@ def get_module_class_from_name(module, name):
 
 def check_model_config(cfg: DictDefault, model_config: Union[AutoConfig, DictDefault]):
     if cfg.is_multimodal:
-        model_config = model_config.text_config
+        if hasattr(model_config, "text_config"):
+            model_config = model_config.text_config
+            model_config.use_cache = False
+        elif hasattr(model_config, "get_text_config"):
+            model_config = model_config.get_text_config()
+            model_config.use_cache = False
+
+        # check if image_size is not set and load image size from model config if available
+        if (
+            cfg.image_size is None
+            and hasattr(model_config, "vision_config")
+            and hasattr(model_config.vision_config, "image_size")
+        ):
+            cfg.image_size = model_config.vision_config.image_size
+            LOG.debug(f"Loaded image size: {cfg.image_size} from model config")
 
     quant_config_exists = (
         hasattr(model_config, "quantization_config")
@@ -440,6 +462,31 @@ def load_processor(cfg: DictDefault, tokenizer: PreTrainedTokenizerBase):
         **processor_kwargs,
     )
 
+    # Attempt to load image size from processor if available
+    if (
+        cfg.image_size is None
+        and hasattr(processor, "size")
+        and any(dim in processor.size for dim in ["width", "height"])
+    ):
+        im_width = None
+        im_height = None
+        if "width" in processor.size:
+            im_width = processor.size["width"]
+        if "height" in processor.size:
+            im_height = processor.size["height"]
+
+        # If both width and height are set, use a tuple
+        if im_width is not None and im_height is not None:
+            cfg.image_size = (im_width, im_height)
+        # If only width is set, use as integer
+        elif im_width is not None:
+            cfg.image_size = im_width
+        # If only height is set, use as integer
+        elif im_height is not None:
+            cfg.image_size = im_height
+
+        LOG.debug(f"Loaded image size: {cfg.image_size} from processor")
+
     return processor
 
 
@@ -477,7 +524,11 @@ class ModelLoader:
         # init model config
         self.model_config = load_model_config(cfg)
         if cfg.is_multimodal:
-            self.text_model_config = self.model_config.text_config
+            if hasattr(self.model_config, "text_config"):
+                self.text_model_config = self.model_config.text_config
+            else:
+                # for qwen2_vl
+                self.text_model_config = self.model_config.get_text_config()
         else:
             self.text_model_config = self.model_config
 
@@ -673,7 +724,7 @@ class ModelLoader:
         should be set according to the type of the model.
         """
         if self.cfg.is_multimodal:
-            self.auto_model_loader = MULTIMODEL_AUTO_MODEL_MAPPING.get(
+            self.auto_model_loader = MULTIMODAL_AUTO_MODEL_MAPPING.get(
                 self.model_config.model_type, AutoModelForVision2Seq
             )
 
@@ -1194,7 +1245,9 @@ class ModelLoader:
             )
         ):
             resize_kwargs = {}
-            if self.cfg.mean_resizing_embeddings is not None:
+            if self.cfg.mean_resizing_embeddings is not None and not (
+                self.model_config.model_type == "llava"
+            ):
                 resize_kwargs["mean_resizing"] = self.cfg.mean_resizing_embeddings
             self.model.resize_token_embeddings(embeddings_len, **resize_kwargs)
         else:

src/axolotl/utils/schemas/config.py

@@ -42,6 +42,7 @@ from axolotl.utils.schemas.model import (
     ModelOutputConfig,
     SpecialTokensConfig,
 )
+from axolotl.utils.schemas.multimodal import MultiModalConfig
 from axolotl.utils.schemas.peft import LoraConfig, ReLoRAConfig
 from axolotl.utils.schemas.training import HyperparametersConfig
 from axolotl.utils.schemas.trl import TRLConfig
@@ -64,6 +65,7 @@ class AxolotlInputConfig(
     LISAConfig,
     GradioConfig,
     RayConfig,
+    MultiModalConfig,
     RemappedParameters,
     DeprecatedParameters,
     BaseModel,

src/axolotl/utils/schemas/enums.py

@@ -22,6 +22,7 @@ class ChatTemplate(str, Enum):
     mistral_v1 = "mistral_v1"  # pylint: disable=invalid-name
     mistral_v2v3 = "mistral_v2v3"  # pylint: disable=invalid-name
     mistral_v3_tekken = "mistral_v3_tekken"  # pylint: disable=invalid-name
+    mistral_v7_tekken = "mistral_v7_tekken"  # pylint: disable=invalid-name
     gemma = "gemma"  # pylint: disable=invalid-name
     cohere = "cohere"  # pylint: disable=invalid-name
     llama3 = "llama3"  # pylint: disable=invalid-name
@@ -36,6 +37,10 @@ class ChatTemplate(str, Enum):
     tokenizer_default = "tokenizer_default"  # pylint: disable=invalid-name
     exaone = "exaone"  # pylint: disable=invalid-name
     metharme = "metharme"  # pylint: disable=invalid-name
+    pixtral = "pixtral"  # pylint: disable=invalid-name
+    llava = "llava"  # pylint: disable=invalid-name
+    qwen2_vl = "qwen2_vl"  # pylint: disable=invalid-name
+    gemma3 = "gemma3"  # pylint: disable=invalid-name
 
 
 class CustomSupportedOptimizers(str, Enum):

src/axolotl/utils/schemas/multimodal.py

@@ -0,0 +1,48 @@
+"""Pydantic models for multimodal-related configuration"""
+
+from typing import Literal
+
+from PIL.Image import Resampling
+from pydantic import BaseModel, Field, field_validator
+
+
+class MultiModalConfig(BaseModel):
+    """Multi-modal configuration subset"""
+
+    image_size: int | tuple[int, int] | None = Field(
+        default=None,
+        json_schema_extra={
+            "description": (
+                "The size of the image to resize to. It can be an integer (resized into padded-square image) or a tuple (width, height)."
+                "If not provided, we will attempt to load from preprocessor.size, otherwise, images won't be resized."
+            )
+        },
+    )
+    image_resize_algorithm: (
+        Literal["bilinear", "bicubic", "lanczos"] | Resampling | None
+    ) = Field(
+        default=None,
+        json_schema_extra={
+            "description": "The resampling algorithm to use for image resizing. Default is bilinear. Please refer to PIL.Image.Resampling for more details."
+        },
+    )
+
+    @field_validator("image_resize_algorithm", mode="before")
+    @classmethod
+    def convert_image_resize_algorithm(cls, image_resize_algorithm):
+        """
+        Convert the image resize algorithm to a PIL.Image.Resampling enum.
+        """
+        if isinstance(image_resize_algorithm, str):
+            image_resize_algorithm = image_resize_algorithm.lower()
+            if image_resize_algorithm == "bilinear":
+                image_resize_algorithm = Resampling.BILINEAR
+            elif image_resize_algorithm == "bicubic":
+                image_resize_algorithm = Resampling.BICUBIC
+            elif image_resize_algorithm == "lanczos":
+                image_resize_algorithm = Resampling.LANCZOS
+            else:
+                raise ValueError(
+                    f"Invalid image resize algorithm: {image_resize_algorithm}"
+                )
+        return image_resize_algorithm

tests/e2e/patched/lora_kernels/test_lora_kernel_patching.py

@@ -144,7 +144,7 @@ def test_swiglu_mlp_integration(small_llama_model):
 def test_geglu_model_integration():
     """Test GeGLU activation with Gemma model."""
     model = AutoModelForCausalLM.from_pretrained(
-        "mhenrichsen/gemma-2b", torch_dtype=torch.float16, device_map="cuda"
+        "mhenrichsen/gemma-2b", torch_dtype=torch.float16, device_map="auto"
     )
     peft_config = get_peft_config(
         {
@@ -347,7 +347,7 @@ def test_model_architecture(model_config):
     """Test LoRA kernel patches across different model architectures."""
     # Load model with appropriate dtype
     model = AutoModelForCausalLM.from_pretrained(
-        model_config["name"], torch_dtype=model_config["dtype"], device_map="cuda"
+        model_config["name"], torch_dtype=model_config["dtype"], device_map="auto"
     )
 
     # Apply LoRA configuration
@@ -408,7 +408,7 @@ def test_kernel_training_integration():
     )
 
     # Load model
-    model, _ = load_model_and_tokenizer(cfg=cfg)
+    model, _, _ = load_model_and_tokenizer(cfg=cfg)
 
     # Verify correct activation function
     layer = model.model.model.layers[0]

tests/e2e/patched/test_sp.py

@@ -8,14 +8,13 @@ import pytest
 import torch
 from accelerate.state import PartialState
 
+from axolotl.monkeypatch.attention.ring_attn import (
+    get_ring_attn_group,
+    set_ring_attn_group,
+)
+from axolotl.utils.collators.batching import adjust_position_ids_for_slice
 from axolotl.utils.dict import DictDefault
 
-# Use a single patch for ring_flash_attn if it's not available
-ring_flash_attn_mock = MagicMock()
-with patch.dict("sys.modules", {"ring_flash_attn": ring_flash_attn_mock}):
-    from axolotl.monkeypatch.attention.ring_attn import get_ring_attn_group
-    from axolotl.utils.collators.batching import adjust_position_ids_for_slice
-
 
 @pytest.fixture
 def partial_state():
@@ -79,6 +78,22 @@ class TestSequenceParallelHelpers:
 class TestRingAttention:
     """Tests for the ring attention functionality."""
 
+    @patch("torch.distributed.get_rank")
+    @patch("torch.distributed.get_world_size")
+    def test_get_ring_attn_group_no_registration(
+        self, mock_world_size, mock_rank, partial_state
+    ):
+        """Test that get_ring_attn_group returns None when no group has been registered."""
+        # Setup mocks
+        mock_world_size.return_value = 4
+        mock_rank.return_value = 0
+
+        # Get the group without registration
+        group = get_ring_attn_group()
+
+        # Verify that None was returned
+        assert group is None
+
     @patch("torch.distributed.new_group")
     @patch("torch.distributed.get_rank")
     @patch("torch.distributed.get_world_size")
@@ -100,24 +115,11 @@ class TestRingAttention:
         # Verify the number of calls without examining the arguments
         assert mock_new_group.call_count == 2
 
-        # Just verify that new_group was called
+        # Verify that new_group was called
         mock_new_group.assert_called()
 
-    @patch("torch.distributed.get_rank")
+        # Clean up
-    @patch("torch.distributed.get_world_size")
+        set_ring_attn_group(None)
-    def test_get_ring_attn_group_no_registration(
-        self, mock_world_size, mock_rank, partial_state
-    ):
-        """Test that get_ring_attn_group returns None when no group has been registered."""
-        # Setup mocks
-        mock_world_size.return_value = 4
-        mock_rank.return_value = 0
-
-        # Get the group without registration
-        group = get_ring_attn_group()
-
-        # Verify that None was returned
-        assert group is None
 
 
 # Mock a simplified DataCollator test

tests/e2e/test_deepseekv3.py

@@ -1,5 +1,5 @@
 """
-E2E tests for lora llama
+E2E tests for deepseekv3
 """
 
 import logging

tests/e2e/test_gemma2.py

@@ -0,0 +1,133 @@
+"""
+E2E tests for gemma2
+"""
+
+import logging
+import os
+from pathlib import Path
+
+import pytest
+
+from axolotl.cli.args import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+LOG = logging.getLogger("axolotl.tests.e2e")
+os.environ["WANDB_DISABLED"] = "true"
+
+
+class TestGemma2:
+    """
+    Test case for Gemma2 models
+    """
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_lora_gemma2(self, temp_dir, sample_packing):
+        # pylint: disable=duplicate-code
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-2-33M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "drop_system_message": True,
+                        "split": "train[:1%]",
+                    },
+                ],
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "chat_template": "gemma",  # gemma2's template is same as gemma
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 4,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "save_safetensors": True,
+                "bf16": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_fft_gemma2(self, temp_dir, sample_packing):
+        # pylint: disable=duplicate-code
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-2-33M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "split": "train[:1%]",
+                        "drop_system_message": True,
+                    },
+                ],
+                "chat_template": "gemma",  # gemma2's template is same as gemma
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 4,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "save_safetensors": True,
+                "bf16": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/test_gemma3_text.py

@@ -0,0 +1,131 @@
+"""
+E2E tests for gemma3_text
+"""
+
+import logging
+import os
+from pathlib import Path
+
+import pytest
+
+from axolotl.cli.args import TrainerCliArgs
+from axolotl.common.datasets import load_datasets
+from axolotl.train import train
+from axolotl.utils.config import normalize_config, validate_config
+from axolotl.utils.dict import DictDefault
+
+LOG = logging.getLogger("axolotl.tests.e2e")
+os.environ["WANDB_DISABLED"] = "true"
+
+
+class TestGemma3Text:
+    """
+    Test case for Gemma3Text models
+    """
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_lora_gemma3_text(self, temp_dir, sample_packing):
+        # pylint: disable=duplicate-code
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-3-34M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "adapter": "lora",
+                "lora_r": 8,
+                "lora_alpha": 16,
+                "lora_dropout": 0.05,
+                "lora_target_linear": True,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "split": "train[:1%]",
+                    },
+                ],
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "chat_template": "gemma3",
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 4,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "save_safetensors": True,
+                "bf16": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "adapter_model.safetensors").exists()
+
+    @pytest.mark.parametrize(
+        "sample_packing",
+        [True, False],
+    )
+    def test_fft_gemma3_text(self, temp_dir, sample_packing):
+        # pylint: disable=duplicate-code
+        cfg = DictDefault(
+            {
+                "base_model": "axolotl-ai-co/gemma-3-34M",
+                "trust_remote_code": True,
+                "sample_packing": sample_packing,
+                "flash_attention": True,
+                "sequence_len": 2048,
+                "val_set_size": 0,
+                "datasets": [
+                    {
+                        "path": "mlabonne/FineTome-100k",
+                        "type": "chat_template",
+                        "field_messages": "conversations",
+                        "message_property_mappings": {
+                            "role": "from",
+                            "content": "value",
+                        },
+                        "split": "train[:1%]",
+                    },
+                ],
+                "chat_template": "gemma3",
+                "special_tokens": {
+                    "bos_token": "<bos>",
+                    "eos_token": "<eos>",
+                },
+                "num_epochs": 1,
+                "micro_batch_size": 1,
+                "gradient_accumulation_steps": 4,
+                "output_dir": temp_dir,
+                "learning_rate": 0.00001,
+                "optimizer": "adamw_bnb_8bit",
+                "lr_scheduler": "cosine",
+                "max_steps": 5,
+                "save_safetensors": True,
+                "bf16": True,
+            }
+        )
+        cfg = validate_config(cfg)
+        normalize_config(cfg)
+        cli_args = TrainerCliArgs()
+        dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)
+
+        train(cfg=cfg, dataset_meta=dataset_meta)
+        assert (Path(temp_dir) / "model.safetensors").exists()

tests/e2e/test_schedulers.py

@@ -54,7 +54,7 @@ class TestCustomSchedulers(unittest.TestCase):
                 "gradient_accumulation_steps": 1,
                 "output_dir": temp_dir,
                 "learning_rate": 0.00001,
-                "optimizer": "adamw_hf",
+                "optimizer": "adamw_torch_fused",
                 "max_steps": 20,
                 "lr_scheduler": "rex",
                 "warmup_steps": 5,