diff --git a/docs/source/grpo_trainer.md b/docs/source/grpo_trainer.md
index ed811769b16..7ba79e92f4a 100644
--- a/docs/source/grpo_trainer.md
+++ b/docs/source/grpo_trainer.md
@@ -228,7 +228,7 @@ training_args = GRPOConfig(
 
 Depending on the model size and the overall GPU memory requirements for training, you may need to adjust the `vllm_gpu_memory_utilization` parameter in [`GRPOConfig`] to avoid underutilization or out-of-memory errors.
 
-We provide a [small script](https://huggingface.co/spaces/trl-lib/recommend-vllm-memory) to help estimate the recommended GPU memory utilization based on your model configuration and experiment settings. Simply use it as follows to get `vllm_gpu_memory_utilization` recommendation:
+We provide a [HF Space](https://huggingface.co/spaces/trl-lib/recommend-vllm-memory) to help estimate the recommended GPU memory utilization based on your model configuration and experiment settings. Simply use it as follows to get `vllm_gpu_memory_utilization` recommendation:
 
 <iframe
 	src="https://trl-lib-recommend-vllm-memory.hf.space"
@@ -242,7 +242,9 @@ If the recommended value does not work in your environment, we suggest adding a
 </Tip>
 
 <Tip>
+
 By default, GRPO uses `MASTER_ADDR=localhost` and `MASTER_PORT=12345` for vLLM, but you can override these values by setting the environment variables accordingly.
+
 </Tip>
 
 For more information, see [Speeding up training with vLLM](speeding_up_training#vllm-for-fast-generation-in-online-methods).
@@ -329,7 +331,7 @@ The [`GRPOTrainer`] supports using custom reward functions instead of dense rewa
      - `prompts` (contains the prompts),
      - `completions` (contains the generated completions),
      - `completions_ids` (contains the tokenized completions),
-     - `trainer_state` ([`transformers.TrainerState`]): The current state of the trainer. This can be used to implement dynamic reward functions, such as curriculum learning, where the reward is adjusted based on the training progress. For more details on the available attributes, refer to the [`TrainerState`](https://huggingface.co/docs/transformers/main/en/main_classes/callback#transformers.TrainerState) documentation.
+     - `trainer_state` ([`~transformers.TrainerState`]): The current state of the trainer. This can be used to implement dynamic reward functions, such as curriculum learning, where the reward is adjusted based on the training progress.
      - All columns names (but `prompt`) that the dataset may have. For example, if the dataset contains a column named `ground_truth`, the function will be called with `ground_truth` as a keyword argument.
 
      The easiest way to comply with this requirement is to use `**kwargs` in the function signature.
@@ -520,10 +522,66 @@ trainer = GRPOTrainer(
     ...,
 )
 ```
+
 and the reward will be computed as the sum of the rewards from each function, or the weighted sum if `reward_weights` is provided in the config.
 
 Note that [`GRPOTrainer`] supports multiple reward functions of different types. See the parameters documentation for more details.
 
+## Vision-Language Model (VLM) Training
+
+GRPO supports training Vision-Language Models (VLMs) on multimodal datasets containing both text and images.
+
+### Supported Models
+
+Tested with:
+
+- **Qwen2.5-VL** — e.g., `Qwen/Qwen2.5-VL-3B-Instruct`
+- **Qwen2-VL** — e.g., `Qwen/Qwen2-VL-2B-Instruct`
+- **Gemma3** — e.g., `google/gemma-3-4b-it`
+  
+<Tip>
+Compatibility with all VLMs is not guaranteed. If you believe a model should be supported, feel free to open an issue on GitHub — or better yet, submit a pull request with the required changes.
+</Tip>
+
+### Quick Start
+
+Use [grpo\_vlm.py](https://github.com/huggingface/trl/blob/main/examples/scripts/grpo_vlm.py) to fine-tune a VLM. Example command for training on [`lmms-lab/multimodal-open-r1-8k-verified`](https://huggingface.co/datasets/lmms-lab/multimodal-open-r1-8k-verified):
+
+```bash
+accelerate launch \
+  --config_file=examples/accelerate_configs/deepspeed_zero3.yaml \
+  examples/scripts/grpo_vlm.py \
+  --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
+  --output_dir grpo-Qwen2.5-VL-3B-Instruct \
+  --learning_rate 1e-5 \
+  --gradient_checkpointing \
+  --torch_dtype bfloat16 \
+  --max_prompt_length 2048 \
+  --max_completion_length 1024 \
+  --use_vllm \
+  --vllm_mode colocate \
+  --use_peft \
+  --lora_target_modules "q_proj", "v_proj" \
+  --log_completions
+```
+
+### Configuration Tips
+
+- Use LoRA on vision-language projection layers
+- VLM may require a lot of image tokens, which cannot be truncated. Set `max_prompt_length` to a higher value (e.g., 2048 in the above example) to accommodate longer prompts.
+- Enable 4-bit quantization to reduce memory usage
+- VLMs are memory-intensive — start with smaller batch sizes
+- Most models are compatible with vLLM (`server` and `colocate` modes)
+
+### Dataset Format
+
+Each training sample should include:
+
+- `prompt`: Text formatted via the processor's chat template
+- `image`: A single image (PIL or NumPy array)
+
+The trainer automatically handles image-to-tensor conversion via the model’s image processor.
+
 ## GRPOTrainer
 
 [[autodoc]] GRPOTrainer
diff --git a/examples/scripts/grpo_vlm.py b/examples/scripts/grpo_vlm.py
new file mode 100644
index 00000000000..66000cc1308
--- /dev/null
+++ b/examples/scripts/grpo_vlm.py
@@ -0,0 +1,168 @@
+# Copyright 2020-2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+pip install math_verify
+
+accelerate launch \
+  --config_file=examples/accelerate_configs/deepspeed_zero3.yaml \
+  examples/scripts/grpo_vlm.py \
+  --model_name_or_path Qwen/Qwen2.5-VL-3B-Instruct \
+  --output_dir grpo-Qwen2.5-VL-3B-Instruct \
+  --learning_rate 1e-5 \
+  --gradient_checkpointing \
+  --torch_dtype bfloat16 \
+  --max_prompt_length 2048 \
+  --max_completion_length 1024 \
+  --use_vllm \
+  --vllm_mode colocate \
+  --use_peft \
+  --lora_target_modules "q_proj", "v_proj" \
+  --log_completions
+"""
+
+import torch
+from datasets import load_dataset
+from latex2sympy2_extended import NormalizationConfig
+from math_verify import LatexExtractionConfig, parse, verify
+
+from trl import (
+    GRPOConfig,
+    GRPOTrainer,
+    ModelConfig,
+    ScriptArguments,
+    TrlParser,
+    get_kbit_device_map,
+    get_peft_config,
+    get_quantization_config,
+)
+from trl.rewards import think_format_reward
+
+
+if __name__ == "__main__":
+    parser = TrlParser((ScriptArguments, GRPOConfig, ModelConfig))
+    script_args, training_args, model_args = parser.parse_args_and_config()
+    ################
+    # Model & Processor
+    ################
+    torch_dtype = (
+        model_args.torch_dtype if model_args.torch_dtype in ["auto", None] else getattr(torch, model_args.torch_dtype)
+    )
+    quantization_config = get_quantization_config(model_args)
+    model_kwargs = dict(
+        revision=model_args.model_revision,
+        attn_implementation=model_args.attn_implementation,
+        torch_dtype=torch_dtype,
+        device_map=get_kbit_device_map() if quantization_config is not None else None,
+        quantization_config=quantization_config,
+    )
+
+    ################
+    # Dataset
+    ################
+    dataset = load_dataset("lmms-lab/multimodal-open-r1-8k-verified", split="train")
+    dataset = dataset.train_test_split(test_size=100, seed=42)
+
+    SYSTEM_PROMPT = (
+        "A conversation between user and assistant. The user asks a question, and the assistant solves it. The "
+        "assistant first thinks about the reasoning process in the mind and then provides the user with the answer. "
+        "The reasoning process and answer are enclosed within <think></think> tags, i.e., <think>\nThis is my "
+        "reasoning.\n</think>\nThis is my answer."
+    )
+
+    def make_conversation(example):
+        prompt = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": example["problem"]},
+        ]
+        return {"prompt": prompt}
+
+    dataset = dataset.map(make_conversation)
+
+    # Filter have big images
+    def filter_big_images(example):
+        image = example["image"]
+        return image.size[0] < 512 and image.size[1] < 512
+
+    dataset = dataset.filter(filter_big_images)
+
+    train_dataset = dataset["train"]
+    eval_dataset = dataset["test"] if training_args.eval_strategy != "no" else None
+
+    ################
+    # Reward Function for Training
+    ################
+    def accuracy_reward(completions, solution: list[str], **kwargs):
+        """Reward function that checks if the completion matches the ground truth.
+        - If both gold and prediction are parseable → use math verification.
+        - If not parseable → compare as normalized text.
+        """
+        rewards = []
+        contents = [completion[0]["content"] for completion in completions]
+        for content, sol in zip(contents, solution):
+            try:
+                gold_parsed = parse(sol, extraction_mode="first_match")
+            except Exception:
+                gold_parsed = []
+
+            if len(gold_parsed) != 0:
+                # Try parsing predicted answer too
+                try:
+                    answer_parsed = parse(
+                        content,
+                        extraction_config=[
+                            LatexExtractionConfig(
+                                normalization_config=NormalizationConfig(
+                                    nits=False,
+                                    malformed_operators=False,
+                                    basic_latex=True,
+                                    boxed="all",
+                                    units=True,
+                                ),
+                                boxed_match_priority=0,
+                                try_extract_without_anchor=False,
+                            )
+                        ],
+                        extraction_mode="first_match",
+                    )
+                    reward = float(verify(gold_parsed, answer_parsed))
+                except Exception as e:
+                    print(f"verify failed: {e}, answer: {content}, gold: {sol}")
+                    reward = None
+            else:
+                # fallback to text match
+                reward = float(content.strip().lower() == sol.strip().lower())
+
+            rewards.append(reward)
+
+        return rewards
+
+    ################
+    # Training
+    ################
+    trainer = GRPOTrainer(
+        model=model_args.model_name_or_path,
+        args=training_args,
+        reward_funcs=[think_format_reward, accuracy_reward],
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        peft_config=get_peft_config(model_args),
+    )
+
+    trainer.train()
+
+    # Save and push to hub
+    trainer.save_model(training_args.output_dir)
+    if training_args.push_to_hub:
+        trainer.push_to_hub(dataset_name=script_args.dataset_name)
diff --git a/scripts/generate_tiny_models.py b/scripts/generate_tiny_models.py
index 9e1d87ea6a4..0b472c9c030 100644
--- a/scripts/generate_tiny_models.py
+++ b/scripts/generate_tiny_models.py
@@ -24,7 +24,6 @@
     BartModel,
     BloomConfig,
     BloomForCausalLM,
-    CLIPVisionConfig,
     CohereConfig,
     CohereForCausalLM,
     DbrxConfig,
@@ -35,6 +34,8 @@
     FalconMambaForCausalLM,
     Gemma2Config,
     Gemma2ForCausalLM,
+    Gemma3Config,
+    Gemma3ForConditionalGeneration,
     GemmaConfig,
     GemmaForCausalLM,
     GPT2Config,
@@ -58,19 +59,21 @@
     PaliGemmaForConditionalGeneration,
     Phi3Config,
     Phi3ForCausalLM,
+    Qwen2_5_VLConfig,
+    Qwen2_5_VLForConditionalGeneration,
     Qwen2Config,
     Qwen2ForCausalLM,
     Qwen2ForSequenceClassification,
+    Qwen2VLConfig,
+    Qwen2VLForConditionalGeneration,
     Qwen3Config,
     Qwen3ForCausalLM,
     Qwen3ForSequenceClassification,
     Qwen3MoeConfig,
     Qwen3MoeForCausalLM,
-    SiglipVisionConfig,
     T5Config,
     T5ForConditionalGeneration,
 )
-from transformers.models.idefics2.configuration_idefics2 import Idefics2VisionConfig
 
 
 ORGANIZATION = "trl-internal-testing"
@@ -211,8 +214,8 @@ def push_to_hub(model, tokenizer, prefix=None, suffix=None):
 
 # Encoder-decoder models
 for model_id, config_class, model_class, suffix in [
-    ("google/flan-t5-small", T5Config, T5ForConditionalGeneration, None),
     ("facebook/bart-base", BartConfig, BartModel, None),
+    ("google/flan-t5-small", T5Config, T5ForConditionalGeneration, None),
 ]:
     tokenizer = AutoTokenizer.from_pretrained(model_id)
     config = config_class(
@@ -232,35 +235,44 @@ def push_to_hub(model, tokenizer, prefix=None, suffix=None):
 
 
 # Vision Language Models
-# fmt: off
-for model_id, config_class, text_config_class, vision_config_class, model_class in [
-    ("HuggingFaceM4/idefics2-8b", Idefics2Config, MistralConfig, Idefics2VisionConfig, Idefics2ForConditionalGeneration),
-    ("llava-hf/llava-1.5-7b-hf", LlavaConfig, LlamaConfig, CLIPVisionConfig, LlavaForConditionalGeneration),
-    ("llava-hf/llava-v1.6-mistral-7b-hf", LlavaNextConfig, MistralConfig, CLIPVisionConfig, LlavaNextForConditionalGeneration),
-    ("google/paligemma-3b-pt-224", PaliGemmaConfig, GemmaConfig, SiglipVisionConfig, PaliGemmaForConditionalGeneration),
+for model_id, config_class, model_class in [
+    ("google/gemma-3-4b-it", Gemma3Config, Gemma3ForConditionalGeneration),
+    ("google/paligemma-3b-pt-224", PaliGemmaConfig, PaliGemmaForConditionalGeneration),
+    ("HuggingFaceM4/idefics2-8b", Idefics2Config, Idefics2ForConditionalGeneration),
+    ("llava-hf/llava-1.5-7b-hf", LlavaConfig, LlavaForConditionalGeneration),
+    ("llava-hf/llava-v1.6-mistral-7b-hf", LlavaNextConfig, LlavaNextForConditionalGeneration),
+    ("Qwen/Qwen2-VL-2B-Instruct", Qwen2VLConfig, Qwen2VLForConditionalGeneration),
+    ("Qwen/Qwen2.5-VL-3B-Instruct", Qwen2_5_VLConfig, Qwen2_5_VLForConditionalGeneration),
 ]:
-# fmt: on
     processor = AutoProcessor.from_pretrained(model_id)
     kwargs = {}
+    text_kwargs = {}
+    vision_kwargs = {}
     if config_class == PaliGemmaConfig:
         kwargs["projection_dim"] = 8
-    vision_kwargs = {}
-    if vision_config_class in [CLIPVisionConfig, SiglipVisionConfig]:
+    if config_class in [LlavaConfig, LlavaNextConfig, PaliGemmaConfig]:
         vision_kwargs["projection_dim"] = 8
-    if vision_config_class == CLIPVisionConfig:
+    if config_class in [LlavaConfig, LlavaNextConfig]:
         vision_kwargs["image_size"] = 336
         vision_kwargs["patch_size"] = 14
+    if config_class in [Qwen2VLConfig, Qwen2_5_VLConfig]:
+        kwargs["vision_start_token_id"] = 151652
+        text_kwargs["rope_scaling"] = {"type": "mrope", "mrope_section": [1]}
+        vision_kwargs["depth"] = 4
+        vision_kwargs["embed_dim"] = 64
+
     config = config_class(
-        text_config=text_config_class(
+        text_config=dict(
             vocab_size=processor.tokenizer.vocab_size + len(processor.tokenizer.added_tokens_encoder),
             hidden_size=8,
             num_attention_heads=4,
             num_key_value_heads=2,
             num_hidden_layers=2,
             intermediate_size=32,
+            **text_kwargs,
         ),
-        vision_config=vision_config_class(
-            hidden_size=8,
+        vision_config=dict(
+            hidden_size=16,
             num_attention_heads=4,
             num_hidden_layers=2,
             intermediate_size=32,
diff --git a/scripts/generate_zen_image_dataset.py b/scripts/generate_zen_image_dataset.py
new file mode 100644
index 00000000000..5795c8d8928
--- /dev/null
+++ b/scripts/generate_zen_image_dataset.py
@@ -0,0 +1,704 @@
+# Copyright 2020-2025 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+
+import numpy as np
+from datasets import Dataset, Features, Image, Sequence, Value
+from transformers import HfArgumentParser
+
+
+Message = [{"content": Value("string"), "role": Value("string")}]
+
+
+@dataclass
+class ScriptArguments:
+    r"""
+    Arguments for the script.
+
+    Args:
+        test_size (`float`, *optional*, defaults to `0.1`):
+            Fraction of the dataset to include in the test split.
+        push_to_hub (`bool`, *optional*, defaults to `False`):
+            Whether to push the dataset to the Hugging Face Hub.
+        repo_id (`str`, *optional*, defaults to `"trl-internal-testing/zen"`):
+            Hugging Face repository ID to push the dataset to.
+    """
+
+    test_size: float = field(
+        default=0.1,
+        metadata={"help": "Fraction of the dataset to include in the test split."},
+    )
+    push_to_hub: bool = field(
+        default=False,
+        metadata={"help": "Whether to push the dataset to the Hugging Face Hub."},
+    )
+    repo_id: str = field(
+        default="trl-internal-testing/zen-image",
+        metadata={"help": "Hugging Face repository ID to push the dataset to."},
+    )
+
+
+def main(test_size, push_to_hub, repo_id):
+    # fmt: off
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "text": [
+            "Beautiful is better than ugly.",
+            "Explicit is better than implicit.",
+            "Simple is better than complex.",
+            "Complex is better than complicated.",
+            "Flat is better than nested.",
+            "Sparse is better than dense.",
+            "Readability counts.",
+            "Special cases aren't special enough to break the rules.",
+            "Although practicality beats purity.",
+            "Errors should never pass silently.",
+            "Unless explicitly silenced.",
+            "In the face of ambiguity, refuse the temptation to guess.",
+            "There should be one-- and preferably only one --obvious way to do it.",
+            "Although that way may not be obvious at first unless you're Dutch.",
+            "Now is better than never.",
+            "Although never is often better than *right* now.",
+            "If the implementation is hard to explain, it's a bad idea.",
+            "If the implementation is easy to explain, it may be a good idea.",
+            "Namespaces are one honking great idea -- let's do more of those!",
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_language_modeling_dataset = Dataset.from_dict(data, features=Features(text=Value("string"), image=Image()))
+    standard_language_modeling_dataset = standard_language_modeling_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_language_modeling_dataset.push_to_hub(repo_id, config_name="standard_language_modeling")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data =  {
+        "prompt": [
+            "Beautiful is better than",
+            "Explicit is",
+            "Simple is better",
+            "Complex",
+            "Flat is better than",
+            "Sparse is better",
+            "Readability",
+            "Special cases aren't special",
+            "Although practicality beats",
+            "Errors should never",
+            "Unless explicitly",
+            "In the face of ambiguity, refuse",
+            "There should be one-- and preferably",
+            "Although that way may not be obvious at first unless you're",
+            "Now is",
+            "Although never is often",
+            "If the implementation is hard to explain,",
+            "If the implementation is easy",
+            "Namespaces are one honking great",
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_prompt_only_dataset = Dataset.from_dict(data, features=Features(prompt=Value("string"), image=Image()))
+    standard_prompt_only_dataset = standard_prompt_only_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_prompt_only_dataset.push_to_hub(repo_id, config_name="standard_prompt_only")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            "Beautiful is better than",
+            "Explicit is",
+            "Simple is better",
+            "Complex",
+            "Flat is better than",
+            "Sparse is better",
+            "Readability",
+            "Special cases aren't special",
+            "Although practicality beats",
+            "Errors should never",
+            "Unless explicitly",
+            "In the face of ambiguity, refuse",
+            "There should be one-- and preferably",
+            "Although that way may not be obvious at first unless you're",
+            "Now is",
+            "Although never is often",
+            "If the implementation is hard to explain,",
+            "If the implementation is easy",
+            "Namespaces are one honking great",
+        ],
+        "completion": [
+            " ugly.",
+            " better than implicit.",
+            " than complex.",
+            " is better than complicated.",
+            " nested.",
+            " than dense.",
+            " counts.",
+            " enough to break the rules.",
+            " purity.",
+            " pass silently.",
+            " silenced.",
+            " the temptation to guess.",
+            " only one --obvious way to do it.",
+            " Dutch.",
+            " better than never.",
+            " better than *right* now.",
+            " it's a bad idea.",
+            " to explain, it may be a good idea.",
+            " idea -- let's do more of those!",
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_prompt_completion_dataset = Dataset.from_dict(data, features=Features(prompt=Value("string"), completion=Value("string"), image=Image()))
+    standard_prompt_completion_dataset = standard_prompt_completion_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_prompt_completion_dataset.push_to_hub(repo_id, config_name="standard_prompt_completion")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            "Beautiful is better than",
+            "Explicit is",
+            "Simple is better",
+            "Complex",
+            "Flat is better than",
+            "Sparse is better",
+            "Readability",
+            "Special cases aren't special",
+            "Although practicality beats",
+            "Errors should never",
+            "Unless explicitly",
+            "In the face of ambiguity, refuse",
+            "There should be one-- and preferably",
+            "Although that way may not be obvious at first unless you're",
+            "Now is",
+            "Although never is often",
+            "If the implementation is hard to explain,",
+            "If the implementation is easy",
+            "Namespaces are one honking great",
+        ],
+        "chosen": [
+            " ugly.",
+            " better than implicit.",
+            " than complex.",
+            " is better than complicated.",
+            " nested.",
+            " than dense.",
+            " counts.",
+            " enough to break the rules.",
+            " purity.",
+            " pass silently.",
+            " silenced.",
+            " the temptation to guess.",
+            " only one --obvious way to do it.",
+            " Dutch.",
+            " better than never.",
+            " better than *right* now.",
+            " it's a bad idea.",
+            " to explain, it may be a good idea.",
+            " idea -- let's do more of those!",
+        ],
+        "rejected": [
+            " the moon.",
+            " worse than nothing.",
+            " than a long vacation.",
+            " is always the answer.",
+            " chocolate.",
+            " without any context.",
+            " is optional.",
+            " enough to become unicorns.",
+            " reality.",
+            " pass their driving test.",
+            " forgotten.",
+            " the opportunity to laugh.",
+            " two or more confusing methods.",
+            " a time traveler.",
+            " never better.",
+            " not even a possibility.",
+            " it's clearly the best choice.",
+            " it's probably magic.",
+            " watermelon -- let's plant some!",
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_preference_dataset = Dataset.from_dict(data, features=Features(prompt=Value("string"), chosen=Value("string"), rejected=Value("string"), image=Image()))
+    standard_preference_dataset = standard_preference_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_preference_dataset.push_to_hub(repo_id, config_name="standard_preference")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "chosen": [
+            "Beautiful is better than ugly.",
+            "Explicit is better than implicit.",
+            "Simple is better than complex.",
+            "Complex is better than complicated.",
+            "Flat is better than nested.",
+            "Sparse is better than dense.",
+            "Readability counts.",
+            "Special cases aren't special enough to break the rules.",
+            "Although practicality beats purity.",
+            "Errors should never pass silently.",
+            "Unless explicitly silenced.",
+            "In the face of ambiguity, refuse the temptation to guess.",
+            "There should be one-- and preferably only one --obvious way to do it.",
+            "Although that way may not be obvious at first unless you're Dutch.",
+            "Now is better than never.",
+            "Although never is often better than *right* now.",
+            "If the implementation is hard to explain, it's a bad idea.",
+            "If the implementation is easy to explain, it may be a good idea.",
+            "Namespaces are one honking great idea -- let's do more of those!",
+        ],
+        "rejected": [
+            "Beautiful is better than the moon.",
+            "Explicit is worse than nothing.",
+            "Simple is better than a long vacation.",
+            "Complex is always the answer.",
+            "Flat is better than chocolate.",
+            "Sparse is better without any context.",
+            "Readability is optional.",
+            "Special cases aren't special enough to become unicorns.",
+            "Although practicality beats reality.",
+            "Errors should never pass their driving test.",
+            "Unless explicitly forgotten.",
+            "In the face of ambiguity, refuse the opportunity to laugh.",
+            "There should be one-- and preferably two or more confusing methods.",
+            "Although that way may not be obvious at first unless you're a time traveler.",
+            "Now is never better.",
+            "Although never is often not even a possibility.",
+            "If the implementation is hard to explain, it's clearly the best choice.",
+            "If the implementation is easy it's probably magic.",
+            "Namespaces are one honking great watermelon -- let's plant some!",
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_implicit_prompt_preference_dataset = Dataset.from_dict(data, features=Features(chosen=Value("string"), rejected=Value("string"), image=Image()))
+    {'prompt': Value(dtype='string'), 'completions': Sequence(feature=Value(dtype='string'), length=-1), 'labels': Sequence(feature=Value(dtype='bool'), length=-1)}
+    standard_implicit_prompt_preference_dataset = standard_implicit_prompt_preference_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_implicit_prompt_preference_dataset.push_to_hub(repo_id, config_name="standard_implicit_prompt_preference")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            "Beautiful is better than",
+            "Explicit is",
+            "Simple is better",
+            "Complex",
+            "Flat is better than",
+            "Sparse is better",
+            "Readability",
+            "Special cases aren't special",
+            "Although practicality beats",
+            "Errors should never",
+            "Unless explicitly",
+            "In the face of ambiguity, refuse",
+            "There should be one-- and preferably",
+            "Although that way may not be obvious at first unless you're",
+            "Now is",
+            "Although never is often",
+            "If the implementation is hard to explain,",
+            "If the implementation is easy",
+            "Namespaces are one honking great",
+        ],
+        "completion": [
+            " ugly.",
+            " worse than nothing.",
+            " than a long vacation.",
+            " is better than complicated.",
+            " nested.",
+            " without any context.",
+            " counts.",
+            " enough to become unicorns.",
+            " purity.",
+            " pass silently.",
+            " forgotten.",
+            " the temptation to guess.",
+            " only one --obvious way to do it.",
+            " a time traveler.",
+            " better than never.",
+            " not even a possibility.",
+            " it's a bad idea.",
+            " it's probably magic.",
+            " watermelon -- let's plant some!",
+        ],
+        "label": [True, False, False, True, True, False, True, False, True, True, False, True, True, False, True, False, True, False, False],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_unpaired_preference_dataset = Dataset.from_dict(data, features=Features(prompt=Value("string"), completion=Value("string"), label=Value("bool"), image=Image()))
+    standard_unpaired_preference_dataset = standard_unpaired_preference_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_unpaired_preference_dataset.push_to_hub(repo_id, config_name="standard_unpaired_preference")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            "Beautiful is better than",
+            "Explicit is better than",
+            "Simple is better than",
+            "Complex is better than",
+            "Flat is better than",
+            "Sparse is better than",
+            "Readability counts",
+            "Special cases aren't special enough",
+            "Although practicality beats",
+            "Errors should never pass",
+            "In the face of ambiguity, refuse",
+            "There should be one-- and preferably only one --",
+            "Although that way may not be",
+            "Now is better than",
+            "Never is often better than",
+            "If the implementation is hard to explain, it's",
+            "If the implementation is easy to explain, it",
+            "Namespaces are one",
+            "Although practicality sometimes beats purity,",
+        ],
+        "completions":[
+            [", let me think...", " ugly."],
+            [", of course,", " implicit.", " because clarity matters."],
+            ["... let's keep it basic,", " complex."],
+            [" when needed,", " complicated."],
+            [" in terms of structure,", " nested."],
+            ["... especially for readability."],
+            [" especially when others read it."],
+            [", unless...", " they follow the rules."],
+            [" some theoretical elegance,", " purity."],
+            [" silently,", " unless explicitly silenced."],
+            [" the temptation to guess."],
+            [" way to do it,"," but sometimes it's not obvious.", " especially when there's more than one possibility."],
+            [" clear at first,", " it will eventually emerge."],
+            [" later."],
+            [" problematic fixes."],
+            [" likely because it's too complicated."],
+            [" might be a good design."],
+            [" of those great ideas,", " that solve many problems."],
+            [" the code should still aim for balance."],
+        ],
+        "labels": [
+            [False, True],
+            [False, True, False],
+            [False, True],
+            [True, True],
+            [True, False],
+            [True],
+            [False],
+            [True, False],
+            [False, False],
+            [False, False],
+            [True],
+            [True, True, False],
+            [True, True],
+            [False],
+            [True], [False],
+            [False],
+            [True, True],
+            [False]
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    standard_stepwise_supervision_dataset = Dataset.from_dict(data, features=Features(prompt=Value("string"), completions=Sequence(Value("string")), labels=Sequence(Value("bool")), image=Image()))
+    standard_stepwise_supervision_dataset = standard_stepwise_supervision_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        standard_stepwise_supervision_dataset.push_to_hub(repo_id, config_name="standard_stepwise_supervision")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "messages": [
+            [{"role": "user", "content": "What is better than ugly?"}, {"role": "assistant", "content": "Beautiful."},],
+            [{"role": "user", "content": "What is better than implicit?"}, {"role": "assistant", "content": "Explicit."}],
+            [{"role": "user", "content": "What is better than complex?"}, {"role": "assistant", "content": "Simple."}],
+            [{"role": "user", "content": "What is better than complicated?"}, {"role": "assistant", "content": "Complex."}],
+            [{"role": "user", "content": "What is better than nested?"}, {"role": "assistant", "content": "Flat."}],
+            [{"role": "user", "content": "What is better than dense?"}, {"role": "assistant", "content": "Sparse."}],
+            [{"role": "user", "content": "What counts?"}, {"role": "assistant", "content": "Readability."}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}, {"role": "assistant", "content": "No, special cases aren't special enough to break the rules."}],
+            [{"role": "user", "content": "What beats purity?"}, {"role": "assistant", "content": "Practicality."}],
+            [{"role": "user", "content": "What should never pass silently?"}, {"role": "assistant", "content": "Errors."}],
+            [{"role": "user", "content": "When can errors pass silently?"}, {"role": "assistant", "content": "When explicitly silenced."}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}, {"role": "assistant", "content": "Refuse the temptation to guess."}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}, {"role": "assistant", "content": "One, and preferably only one."}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}, {"role": "assistant", "content": "Dutch."}],
+            [{"role": "user", "content": "What is better than never?"}, {"role": "assistant", "content": "Now is better than never."}],
+            [{"role": "user", "content": "Is never better than *right* now?"}, {"role": "assistant", "content": "Yes, often."}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}, {"role": "assistant", "content": "It means it's a bad idea."}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}, {"role": "assistant", "content": "It means it may be a good idea."}],
+            [{"role": "user", "content": "Any great ideas?"}, {"role": "assistant", "content": "Namespaces are one honking great idea."}],
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    conversational_language_modeling_dataset = Dataset.from_dict(data, features=Features(messages=Message, image=Image()))
+    conversational_language_modeling_dataset = conversational_language_modeling_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        conversational_language_modeling_dataset.push_to_hub(repo_id, config_name="conversational_language_modeling")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            [{"role": "user", "content": "What is better than ugly?"}],
+            [{"role": "user", "content": "What is better than implicit?"}],
+            [{"role": "user", "content": "What is better than complex?"}],
+            [{"role": "user", "content": "What is better than complicated?"}],
+            [{"role": "user", "content": "What is better than nested?"}],
+            [{"role": "user", "content": "What is better than dense?"}],
+            [{"role": "user", "content": "What counts?"}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}],
+            [{"role": "user", "content": "What beats purity?"}],
+            [{"role": "user", "content": "What should never pass silently?"}],
+            [{"role": "user", "content": "When can errors pass silently?"}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}],
+            [{"role": "user", "content": "What is better than never?"}],
+            [{"role": "user", "content": "Is never better than *right* now?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}],
+            [{"role": "user", "content": "Any great ideas?"}],
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    conversational_prompt_only_dataset = Dataset.from_dict(data, features=Features(prompt=Message, image=Image()))
+    conversational_prompt_only_dataset = conversational_prompt_only_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        conversational_prompt_only_dataset.push_to_hub(repo_id, config_name="conversational_prompt_only")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            [{"role": "user", "content": "What is better than ugly?"}],
+            [{"role": "user", "content": "What is better than implicit?"}],
+            [{"role": "user", "content": "What is better than complex?"}],
+            [{"role": "user", "content": "What is better than complicated?"}],
+            [{"role": "user", "content": "What is better than nested?"}],
+            [{"role": "user", "content": "What is better than dense?"}],
+            [{"role": "user", "content": "What counts?"}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}],
+            [{"role": "user", "content": "What beats purity?"}],
+            [{"role": "user", "content": "What should never pass silently?"}],
+            [{"role": "user", "content": "When can errors pass silently?"}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}],
+            [{"role": "user", "content": "What is better than never?"}],
+            [{"role": "user", "content": "Is never better than *right* now?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}],
+            [{"role": "user", "content": "Any great ideas?"}],
+        ],
+        "completion": [
+            [{"role": "assistant", "content": "Beautiful."}],
+            [{"role": "assistant", "content": "Explicit."}],
+            [{"role": "assistant", "content": "Simple."}],
+            [{"role": "assistant", "content": "Complex."}],
+            [{"role": "assistant", "content": "Flat."}],
+            [{"role": "assistant", "content": "Sparse."}],
+            [{"role": "assistant", "content": "Readability."}],
+            [{"role": "assistant", "content": "No, special cases aren't special enough to break the rules."}],
+            [{"role": "assistant", "content": "Practicality."}],
+            [{"role": "assistant", "content": "Errors."}],
+            [{"role": "assistant", "content": "When explicitly silenced."}],
+            [{"role": "assistant", "content": "Refuse the temptation to guess."}],
+            [{"role": "assistant", "content": "One, and preferably only one."}],
+            [{"role": "assistant", "content": "Dutch."}],
+            [{"role": "assistant", "content": "Now is better than never."}],
+            [{"role": "assistant", "content": "Yes, often."}],
+            [{"role": "assistant", "content": "It means it's a bad idea."}],
+            [{"role": "assistant", "content": "It means it may be a good idea."}],
+            [{"role": "assistant", "content": "Namespaces are one honking great idea."}],
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    conversational_prompt_completion_dataset = Dataset.from_dict(data, features=Features(prompt=Message, completion=Message, image=Image()))
+    conversational_prompt_completion_dataset = conversational_prompt_completion_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        conversational_prompt_completion_dataset.push_to_hub(repo_id, config_name="conversational_prompt_completion")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            [{"role": "user", "content": "What is better than ugly?"}],
+            [{"role": "user", "content": "What is better than implicit?"}],
+            [{"role": "user", "content": "What is better than complex?"}],
+            [{"role": "user", "content": "What is better than complicated?"}],
+            [{"role": "user", "content": "What is better than nested?"}],
+            [{"role": "user", "content": "What is better than dense?"}],
+            [{"role": "user", "content": "What counts?"}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}],
+            [{"role": "user", "content": "What beats purity?"}],
+            [{"role": "user", "content": "What should never pass silently?"}],
+            [{"role": "user", "content": "When can errors pass silently?"}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}],
+            [{"role": "user", "content": "What is better than never?"}],
+            [{"role": "user", "content": "Is never better than *right* now?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}],
+            [{"role": "user", "content": "Any great ideas?"}],
+        ],
+        "chosen": [
+            [{"role": "assistant", "content": "Beautiful."}],
+            [{"role": "assistant", "content": "Explicit."}],
+            [{"role": "assistant", "content": "Simple."}],
+            [{"role": "assistant", "content": "Complex."}],
+            [{"role": "assistant", "content": "Flat."}],
+            [{"role": "assistant", "content": "Sparse."}],
+            [{"role": "assistant", "content": "Readability."}],
+            [{"role": "assistant", "content": "No, special cases aren't special enough to break the rules."}],
+            [{"role": "assistant", "content": "Practicality."}],
+            [{"role": "assistant", "content": "Errors."}],
+            [{"role": "assistant", "content": "When explicitly silenced."}],
+            [{"role": "assistant", "content": "Refuse the temptation to guess."}],
+            [{"role": "assistant", "content": "One, and preferably only one."}],
+            [{"role": "assistant", "content": "Dutch."}],
+            [{"role": "assistant", "content": "Now is better than never."}],
+            [{"role": "assistant", "content": "Yes, often."}],
+            [{"role": "assistant", "content": "It means it's a bad idea."}],
+            [{"role": "assistant", "content": "It means it may be a good idea."}],
+            [{"role": "assistant", "content": "Namespaces are one honking great idea."}],
+        ],
+        "rejected": [
+            [{"role": "assistant", "content": "Acceptable."}],
+            [{"role": "assistant", "content": "Explained."}],
+            [{"role": "assistant", "content": "Very complex."}],
+            [{"role": "assistant", "content": "Very complicated."}],
+            [{"role": "assistant", "content": "Circular."}],
+            [{"role": "assistant", "content": "Heavy."}],
+            [{"role": "assistant", "content": "Looking complicated."}],
+            [{"role": "assistant", "content": "Yes, special cases are special enough to break the rules."}],
+            [{"role": "assistant", "content": "Nothing."}],
+            [{"role": "assistant", "content": "Warnings."}],
+            [{"role": "assistant", "content": "Never."}],
+            [{"role": "assistant", "content": "Give up."}],
+            [{"role": "assistant", "content": "As many as possible."}],
+            [{"role": "assistant", "content": "French."}],
+            [{"role": "assistant", "content": "Some day."}],
+            [{"role": "assistant", "content": "No, never."}],
+            [{"role": "assistant", "content": "It means it's a good idea."}],
+            [{"role": "assistant", "content": "It means it's a bad idea."}],
+            [{"role": "assistant", "content": "Recursion."}],
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    conversational_preference_dataset = Dataset.from_dict(data, features=Features(prompt=Message, chosen=Message, rejected=Message, image=Image()))
+    conversational_preference_dataset = conversational_preference_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        conversational_preference_dataset.push_to_hub(repo_id, config_name="conversational_preference")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "chosen": [
+            [{"role": "user", "content": "What is better than ugly?"}, {"role": "assistant", "content": "Beautiful."}],
+            [{"role": "user", "content": "What is better than implicit?"}, {"role": "assistant", "content": "Explicit."}],
+            [{"role": "user", "content": "What is better than complex?"}, {"role": "assistant", "content": "Simple."}],
+            [{"role": "user", "content": "What is better than complicated?"}, {"role": "assistant", "content": "Complex."}],
+            [{"role": "user", "content": "What is better than nested?"}, {"role": "assistant", "content": "Flat."}],
+            [{"role": "user", "content": "What is better than dense?"}, {"role": "assistant", "content": "Sparse."}],
+            [{"role": "user", "content": "What counts?"}, {"role": "assistant", "content": "Readability."}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}, {"role": "assistant", "content": "No, special cases aren't special enough to break the rules."}],
+            [{"role": "user", "content": "What beats purity?"}, {"role": "assistant", "content": "Practicality."}],
+            [{"role": "user", "content": "What should never pass silently?"}, {"role": "assistant", "content": "Errors."}],
+            [{"role": "user", "content": "When can errors pass silently?"}, {"role": "assistant", "content": "When explicitly silenced."}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}, {"role": "assistant", "content": "Refuse the temptation to guess."}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}, {"role": "assistant", "content": "One, and preferably only one."}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}, {"role": "assistant", "content": "Dutch."}],
+            [{"role": "user", "content": "What is better than never?"}, {"role": "assistant", "content": "Now is better than never."}],
+            [{"role": "user", "content": "Is never better than *right* now?"}, {"role": "assistant", "content": "Yes, often."}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}, {"role": "assistant", "content": "It means it's a bad idea."}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}, {"role": "assistant", "content": "It means it may be a good idea."}],
+            [{"role": "user", "content": "Any great ideas?"}, {"role": "assistant", "content": "Namespaces are one honking great idea."}],
+        ],
+        "rejected": [
+            [{"role": "user", "content": "What is better than ugly?"}, {"role": "assistant", "content": "Acceptable."}],
+            [{"role": "user", "content": "What is better than implicit?"}, {"role": "assistant", "content": "Explained."}],
+            [{"role": "user", "content": "What is better than complex?"}, {"role": "assistant", "content": "Very complex."}],
+            [{"role": "user", "content": "What is better than complicated?"}, {"role": "assistant", "content": "Very complicated."}],
+            [{"role": "user", "content": "What is better than nested?"}, {"role": "assistant", "content": "Circular."}],
+            [{"role": "user", "content": "What is better than dense?"}, {"role": "assistant", "content": "Heavy."}],
+            [{"role": "user", "content": "What counts?"}, {"role": "assistant", "content": "Looking complicated."}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}, {"role": "assistant", "content": "Yes, special cases are special enough to break the rules."}],
+            [{"role": "user", "content": "What beats purity?"}, {"role": "assistant", "content": "Nothing."}],
+            [{"role": "user", "content": "What should never pass silently?"}, {"role": "assistant", "content": "Warnings."}],
+            [{"role": "user", "content": "When can errors pass silently?"}, {"role": "assistant", "content": "Never."}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}, {"role": "assistant", "content": "Give up."}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}, {"role": "assistant", "content": "As many as possible."}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}, {"role": "assistant", "content": "French."}],
+            [{"role": "user", "content": "What is better than never?"}, {"role": "assistant", "content": "Some day."}],
+            [{"role": "user", "content": "Is never better than *right* now?"}, {"role": "assistant", "content": "No, never."}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}, {"role": "assistant", "content": "It means it's a good idea."}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}, {"role": "assistant", "content": "It means it's a bad idea."}],
+            [{"role": "user", "content": "Any great ideas?"}, {"role": "assistant", "content": "Recursion."}],
+        ],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    conversational_implicit_prompt_preference_dataset = Dataset.from_dict(data, features=Features(chosen=Message, rejected=Message, image=Image()))
+    conversational_implicit_prompt_preference_dataset = conversational_implicit_prompt_preference_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        conversational_implicit_prompt_preference_dataset.push_to_hub(repo_id, config_name="conversational_implicit_prompt_preference")
+
+    sizes = np.random.randint(32, 64, size=(19, 2))
+    data = {
+        "prompt": [
+            [{"role": "user", "content": "What is better than ugly?"}],
+            [{"role": "user", "content": "What is better than implicit?"}],
+            [{"role": "user", "content": "What is better than complex?"}],
+            [{"role": "user", "content": "What is better than complicated?"}],
+            [{"role": "user", "content": "What is better than nested?"}],
+            [{"role": "user", "content": "What is better than dense?"}],
+            [{"role": "user", "content": "What counts?"}],
+            [{"role": "user", "content": "Are special cases enough to break the rules?"}],
+            [{"role": "user", "content": "What beats purity?"}],
+            [{"role": "user", "content": "What should never pass silently?"}],
+            [{"role": "user", "content": "When can errors pass silently?"}],
+            [{"role": "user", "content": "What should you do in the face of ambiguity?"}],
+            [{"role": "user", "content": "How many ways should there be to do it?"}],
+            [{"role": "user", "content": "For whom may the way not be obvious at first?"}],
+            [{"role": "user", "content": "What is better than never?"}],
+            [{"role": "user", "content": "Is never better than *right* now?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is hard to explain?"}],
+            [{"role": "user", "content": "What does it mean if the implementation is easy to explain?"}],
+            [{"role": "user", "content": "Any great ideas?"}],
+        ],
+        "completion": [
+            [{'role': 'assistant', 'content': 'Beautiful.'}],
+            [{'role': 'assistant', 'content': 'Explicit.'}],
+            [{'role': 'assistant', 'content': 'Simple.'}],
+            [{'role': 'assistant', 'content': 'Very complicated.'}],
+            [{'role': 'assistant', 'content': 'Flat.'}],
+            [{'role': 'assistant', 'content': 'Sparse.'}],
+            [{'role': 'assistant', 'content': 'Readability.'}],
+            [{'role': 'assistant', 'content': 'Yes, special cases are special enough to break the rules.'}],
+            [{'role': 'assistant', 'content': 'Practicality.'}],
+            [{'role': 'assistant', 'content': 'Warnings.'}],
+            [{'role': 'assistant', 'content': 'When explicitly silenced.'}],
+            [{'role': 'assistant', 'content': 'Give up.'}],
+            [{'role': 'assistant', 'content': 'One, and preferably only one.'}],
+            [{'role': 'assistant', 'content': 'French.'}],
+            [{'role': 'assistant', 'content': 'Some day.'}],
+            [{'role': 'assistant', 'content': 'Yes, often.'}],
+            [{'role': 'assistant', 'content': "It means it's a bad idea."}],
+            [{'role': 'assistant', 'content': 'It means it may be a good idea.'}],
+            [{'role': 'assistant', 'content': 'Namespaces are one honking great idea.'}],
+        ],
+        "label": [True, True, True, False, True, True, True, False, True, False, True, False, True, False, False, True, True, True, True],
+        "image": [np.random.uniform(low=0.0, high=255.0, size=(h, w, 3)).astype(np.uint8) for h, w in sizes],
+    }
+    conversational_unpaired_preference_dataset = Dataset.from_dict(data, features=Features(prompt=Message, completion=Message, label=Value("bool"), image=Image()))
+    conversational_unpaired_preference_dataset = conversational_unpaired_preference_dataset.train_test_split(test_size=test_size, shuffle=False)
+    if push_to_hub:
+        conversational_unpaired_preference_dataset.push_to_hub(repo_id, config_name="conversational_unpaired_preference")
+    # fmt: on
+
+
+if __name__ == "__main__":
+    parser = HfArgumentParser(ScriptArguments)
+    script_args = parser.parse_args_into_dataclasses()[0]
+    main(script_args.test_size, script_args.push_to_hub, script_args.repo_id)
diff --git a/setup.cfg b/setup.cfg
index 744c7f839d6..5ea3a285e82 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -71,6 +71,7 @@ vllm =
 
 vlm =
     Pillow
+    torchvision
 dev =
     %(bco)s
     %(deepspeed)s
diff --git a/tests/slow/test_grpo_slow.py b/tests/slow/test_grpo_slow.py
index d7c490f4115..d2c604825e5 100644
--- a/tests/slow/test_grpo_slow.py
+++ b/tests/slow/test_grpo_slow.py
@@ -13,28 +13,46 @@
 # limitations under the License.
 
 import gc
+import os
 import tempfile
 import unittest
+import warnings
 
+import numpy as np
 import pytest
 import torch
 from accelerate.utils.memory import release_memory
-from datasets import load_dataset
+from datasets import Dataset, Features, Image, Value, load_dataset
 from parameterized import parameterized
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoProcessor,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+)
 from transformers.testing_utils import (
     backend_empty_cache,
+    require_bitsandbytes,
+    require_flash_attn,
     require_liger_kernel,
     require_peft,
     require_torch_accelerator,
     torch_device,
 )
+from transformers.utils import is_peft_available
 
 from trl import GRPOConfig, GRPOTrainer
+from trl.trainer.utils import get_kbit_device_map
 
+from ..testing_utils import require_vllm
 from .testing_constants import MODELS_TO_TEST
 
 
+if is_peft_available():
+    from peft import LoraConfig, PeftModel
+
+
 @pytest.mark.slow
 @require_torch_accelerator
 class GRPOTrainerSlowTester(unittest.TestCase):
@@ -188,3 +206,357 @@ def test_training_with_transformers_paged(self, model_name):
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
 
         release_memory(model, trainer)
+
+    @require_flash_attn
+    @require_bitsandbytes
+    @require_peft
+    @parameterized.expand(
+        [
+            ("HuggingFaceTB/SmolVLM-Instruct",),  # Only test the smaller model to avoid OOM
+        ]
+    )
+    def test_vlm_training(self, model_name):
+        """
+        Test VLM training with aggressive memory optimization.
+
+        This test uses multiple memory reduction techniques:
+        - 4-bit quantization with double quantization
+        - LoRA with very low rank (r=4)
+        - Minimal batch size (1) with gradient accumulation
+        - Small images (64x64 instead of 224x224)
+        - Short sequences (max_completion_length=8)
+        - Only 4 training samples
+        - Only 1 training step
+        - Gradient checkpointing and bfloat16
+        """
+
+        # Create processor once outside the data generator
+        processor = AutoProcessor.from_pretrained(model_name, use_fast=True, padding_side="left")
+        conversation = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "image"},
+                    {"type": "text", "text": "What is in the image?"},
+                ],
+            },
+        ]
+        prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+        def data_gen(num_samples):
+            for _ in range(num_samples):
+                yield {
+                    "prompt": prompt,
+                    "image": np.random.uniform(low=0.0, high=255.0, size=(64, 64, 3)).astype(
+                        np.uint8
+                    ),  # Much smaller images
+                }
+
+        dataset = Dataset.from_generator(
+            data_gen, gen_kwargs={"num_samples": 4}, features=Features(image=Image(), prompt=Value(dtype="string"))
+        )
+        # reduce memory requirements as much as possible
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype="bfloat16",
+            bnb_4bit_quant_type="nf4",
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_storage="bfloat16",
+        )
+        model = AutoModelForImageTextToText.from_pretrained(
+            model_name,
+            attn_implementation="flash_attention_2",
+            torch_dtype="bfloat16",
+            device_map=get_kbit_device_map(),
+            quantization_config=quantization_config,
+        )
+
+        def reward_func(prompts, completions, **kwargs):
+            # simple nonsensical reward
+            return [-((len(c) - 25) ** 2) + 100 for c in completions]
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,
+                per_device_train_batch_size=1,  # Minimal batch size
+                gradient_accumulation_steps=2,  # Maintain effective batch size
+                num_generations=2,
+                max_completion_length=8,  # Much shorter completions
+                max_prompt_length=None,  # Don't limit prompt length for VLM
+                bf16=True,  # Use bfloat16 precision
+                max_steps=1,  # Only do 1 training step to save time and memory
+                report_to="none",
+                logging_strategy="no",
+            )
+            lora_config = LoraConfig(
+                task_type="CAUSAL_LM",
+                r=4,  # Much lower rank for minimal memory
+                lora_alpha=8,  # Reduced alpha proportionally
+                lora_dropout=0.1,
+                target_modules=["q_proj", "v_proj"],  # Minimal target modules
+                # For VLM models, we typically want to freeze the vision encoder
+                # and only adapt the language model parameters
+                modules_to_save=None,
+            )
+
+            try:
+                trainer = GRPOTrainer(
+                    model=model,
+                    processing_class=processor,
+                    reward_funcs=[reward_func],
+                    args=training_args,
+                    train_dataset=dataset,
+                    peft_config=lora_config,
+                )
+
+                self.assertIsInstance(trainer.model, PeftModel)
+
+                previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+                trainer.train()
+
+                self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+                # Check that LoRA parameters have changed
+                # For VLM models, we're more permissive about which parameters can change
+                lora_params_changed = False
+                for n, param in previous_trainable_params.items():
+                    new_param = trainer.model.get_parameter(n)
+                    if "lora" in n.lower():  # LoRA parameters should change
+                        if not torch.equal(param, new_param):
+                            lora_params_changed = True
+
+                # At least some LoRA parameters should have changed during training
+                self.assertTrue(lora_params_changed, "No LoRA parameters were updated during training.")
+
+            except torch.OutOfMemoryError as e:
+                self.skipTest(f"Skipping VLM training test due to insufficient GPU memory: {e}")
+            except Exception as e:
+                # Check for other memory-related errors
+                if any(keyword in str(e).lower() for keyword in ["memory", "cuda", "out of memory", "insufficient"]):
+                    self.skipTest(f"Skipping VLM training test due to hardware constraints: {e}")
+                else:
+                    raise
+
+        release_memory(model, trainer)
+
+    @require_vllm
+    @require_bitsandbytes
+    @require_peft
+    def test_vlm_processor_vllm_colocate_mode(self):
+        """
+        Test that VLM processors work with vLLM in colocate mode.
+
+        This test uses multiple memory optimization techniques to ensure it runs on limited hardware:
+        - LoRA (Low-Rank Adaptation) with minimal rank (r=4)
+        - 4-bit quantization with BitsAndBytesConfig
+        - Gradient checkpointing
+        - bfloat16 precision
+        - Minimal batch sizes and sequence lengths
+        - Very low GPU memory utilization (5%)
+        """
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            config = GRPOConfig(
+                output_dir=tmp_dir,
+                per_device_train_batch_size=1,  # Minimal batch size
+                gradient_accumulation_steps=2,  # Make effective batch size 2, divisible by num_generations
+                num_generations=2,
+                max_completion_length=4,  # Very short completions to reduce memory
+                max_prompt_length=32,  # Very short prompts to reduce memory
+                use_vllm=True,  # Enable vLLM
+                vllm_mode="colocate",  # Use colocate mode to avoid server dependency
+                vllm_gpu_memory_utilization=0.05,  # Use minimal GPU memory (5%)
+                gradient_checkpointing=True,  # Enable gradient checkpointing to save memory
+                bf16=True,  # Use bfloat16 to reduce memory
+                report_to="none",
+                logging_strategy="no",
+            )
+
+            # Create a VLM processor
+            processor = AutoProcessor.from_pretrained(
+                "HuggingFaceTB/SmolVLM-Instruct", use_fast=True, padding_side="left"
+            )
+
+            # Verify processor has both required attributes for VLM detection
+            self.assertTrue(hasattr(processor, "tokenizer"))
+            self.assertTrue(hasattr(processor, "image_processor"))
+
+            def dummy_reward_func(completions, **kwargs):
+                return [1.0] * len(completions)
+
+            # Use LoRA configuration for memory efficiency
+            lora_config = LoraConfig(
+                r=4,  # Very low rank for minimal memory
+                lora_alpha=8,
+                target_modules=["q_proj", "v_proj"],  # Minimal target modules
+                lora_dropout=0.1,
+                bias="none",
+                task_type="CAUSAL_LM",
+            )
+
+            # Use 4-bit quantization for further memory reduction
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.bfloat16,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True,
+            )
+
+            original_env = {}
+            required_env_vars = {
+                "RANK": "0",
+                "LOCAL_RANK": "0",
+                "WORLD_SIZE": "1",
+                "LOCAL_WORLD_SIZE": "1",
+                "MASTER_ADDR": "localhost",
+                "MASTER_PORT": "12355",
+            }
+
+            for key, value in required_env_vars.items():
+                original_env[key] = os.environ.get(key)
+                os.environ[key] = value
+
+            try:
+                # Test VLM processor with vLLM colocate mode
+                with warnings.catch_warnings(record=True) as w:
+                    warnings.simplefilter("always")
+                    try:
+                        # Load model with quantization for memory efficiency
+                        model = AutoModelForCausalLM.from_pretrained(
+                            "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
+                            quantization_config=quantization_config,
+                            torch_dtype=torch.bfloat16,
+                        )
+
+                        trainer = GRPOTrainer(
+                            model=model,
+                            reward_funcs=dummy_reward_func,
+                            args=config,
+                            train_dataset=dataset,
+                            processing_class=processor,  # VLM processor
+                            peft_config=lora_config,  # Use LoRA for memory efficiency
+                        )
+
+                        # Should detect VLM processor correctly and allow vLLM
+                        self.assertTrue(trainer.use_vllm, "vLLM should be enabled for VLM processors in colocate mode")
+                        self.assertEqual(trainer.vllm_mode, "colocate", "Should use colocate mode")
+
+                        # Check if signature columns were set properly
+                        if trainer._signature_columns is not None:
+                            # Should include 'image' in signature columns for VLM processors
+                            self.assertIn(
+                                "image",
+                                trainer._signature_columns,
+                                "Should include 'image' in signature columns for VLM",
+                            )
+
+                        # Should not emit any warnings about VLM incompatibility
+                        incompatibility_warnings = [
+                            str(w_item.message)
+                            for w_item in w
+                            if "does not support VLMs" in str(w_item.message)
+                            or "not compatible" in str(w_item.message).lower()
+                        ]
+                        self.assertEqual(
+                            len(incompatibility_warnings),
+                            0,
+                            f"Should not emit VLM incompatibility warnings, but got: {incompatibility_warnings}",
+                        )
+
+                        # Test passes if we get this far without exceptions
+
+                    except Exception as e:
+                        # If vLLM fails to initialize due to hardware constraints or other issues, that's expected
+                        if any(
+                            keyword in str(e).lower()
+                            for keyword in [
+                                "outofmemoryerror",
+                                "cuda",
+                                "memory",
+                                "insufficient",
+                                "no such device",
+                                "free memory",
+                                "gpu memory utilization",
+                                "decrease gpu memory",
+                            ]
+                        ):
+                            self.skipTest(f"Skipping vLLM colocate test due to hardware constraints: {e}")
+                        elif "KeyError" in str(e) and "RANK" in str(e):
+                            self.skipTest(f"Skipping vLLM colocate test due to environment setup issues: {e}")
+                        elif "ValueError" in str(e) and "memory" in str(e).lower():
+                            self.skipTest(f"Skipping vLLM colocate test due to memory constraints: {e}")
+                        else:
+                            raise
+            finally:
+                # Restore original environment variables
+                for key, original_value in original_env.items():
+                    if original_value is None:
+                        os.environ.pop(key, None)
+                    else:
+                        os.environ[key] = original_value
+
+                release_memory(model, trainer)
+
+    @require_vllm
+    def test_training_vllm(self):
+        """Test that training works with vLLM for generation."""
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                report_to="none",
+                logging_strategy="no",
+                use_vllm=True,
+            )
+
+            try:
+                trainer = GRPOTrainer(
+                    model="Qwen/Qwen2.5-0.5B-Instruct",  # tiny models are too small for vLLM
+                    reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                    args=training_args,
+                    train_dataset=dataset,
+                )
+
+                previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+                trainer.train()
+
+                self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+                # Check that the params have changed
+                for n, param in previous_trainable_params.items():
+                    new_param = trainer.model.get_parameter(n)
+                    self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+            except Exception as e:
+                # If vLLM fails to initialize due to hardware constraints or other issues, that's expected
+                if any(
+                    keyword in str(e).lower()
+                    for keyword in [
+                        "outofmemoryerror",
+                        "cuda",
+                        "memory",
+                        "insufficient",
+                        "no such device",
+                        "free memory",
+                        "gpu memory utilization",
+                        "decrease gpu memory",
+                    ]
+                ):
+                    self.skipTest(f"Skipping vLLM training test due to hardware constraints: {e}")
+                elif "KeyError" in str(e) and "RANK" in str(e):
+                    self.skipTest(f"Skipping vLLM training test due to environment setup issues: {e}")
+                elif "ValueError" in str(e) and "memory" in str(e).lower():
+                    self.skipTest(f"Skipping vLLM training test due to memory constraints: {e}")
+                else:
+                    raise
+
+        release_memory(trainer.model, trainer)
diff --git a/tests/test_grpo_trainer.py b/tests/test_grpo_trainer.py
index 7dbc112077a..9084598ab0c 100644
--- a/tests/test_grpo_trainer.py
+++ b/tests/test_grpo_trainer.py
@@ -19,12 +19,25 @@
 import torch
 from datasets import load_dataset
 from parameterized import parameterized
-from transformers import AutoModelForCausalLM, AutoModelForSequenceClassification, AutoTokenizer
-from transformers.testing_utils import require_peft
+from transformers import (
+    AutoModelForCausalLM,
+    AutoModelForImageTextToText,
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+)
+from transformers.testing_utils import require_liger_kernel, require_peft, require_vision
 from transformers.utils import is_peft_available
 
 from trl import GRPOConfig, GRPOTrainer
-from trl.trainer.grpo_trainer import RepeatSampler, get_high_entropy_mask, shuffle_tensor_dict, split_tensor_dict
+from trl.trainer.grpo_trainer import (
+    RepeatSampler,
+    get_high_entropy_mask,
+    shuffle_sequence_dict,
+    split_pixel_values_by_grid,
+    split_tensor_dict,
+    truncate_with_protected_tokens,
+    unsplit_pixel_values_by_grid,
+)
 
 from .testing_utils import require_vllm
 
@@ -61,13 +74,13 @@ def test_with_none_tensor(self):
             self.assertIsNone(result[i]["y"])
 
 
-class ShuffleTensorDictTester(unittest.TestCase):
+class ShuffleSequenceDictTester(unittest.TestCase):
     def test_shuffle_preserves_shape(self):
         x = torch.arange(6).reshape(3, 2)
         y = torch.arange(3).reshape(3, 1)
         tensor_dict = {"x": x.clone(), "y": y.clone()}
 
-        shuffled = shuffle_tensor_dict(tensor_dict)
+        shuffled = shuffle_sequence_dict(tensor_dict)
 
         self.assertEqual(shuffled["x"].shape, x.shape)
         self.assertEqual(shuffled["y"].shape, y.shape)
@@ -78,7 +91,7 @@ def test_shuffle_consistent_across_tensors(self):
         y = torch.tensor([[1], [2], [3]])
         tensor_dict = {"x": x.clone(), "y": y.clone()}
 
-        shuffled = shuffle_tensor_dict(tensor_dict)
+        shuffled = shuffle_sequence_dict(tensor_dict)
 
         # Build a reverse map from shuffled x rows to y values
         for i in range(3):
@@ -98,11 +111,33 @@ def test_none_tensor_remains_none(self):
         x = torch.arange(6).reshape(3, 2)
         tensor_dict = {"x": x.clone(), "y": None}
 
-        shuffled = shuffle_tensor_dict(tensor_dict)
+        shuffled = shuffle_sequence_dict(tensor_dict)
 
         self.assertIsNone(shuffled["y"])
         self.assertEqual(shuffled["x"].shape, x.shape)
 
+    def test_shuffle_with_list(self):
+        x = torch.tensor([[10, 11], [20, 21], [30, 31]])
+        y = ["a", "b", "c"]
+
+        sequence_dict = {"x": x.clone(), "y": y}
+
+        shuffled = shuffle_sequence_dict(sequence_dict)
+
+        # Check that the list y is shuffled in the same order as x
+        for i in range(3):
+            x_row = shuffled["x"][i]
+            y_val = shuffled["y"][i]
+
+            if torch.equal(x_row, torch.tensor([10, 11])):
+                self.assertEqual(y_val, "a")
+            elif torch.equal(x_row, torch.tensor([20, 21])):
+                self.assertEqual(y_val, "b")
+            elif torch.equal(x_row, torch.tensor([30, 31])):
+                self.assertEqual(y_val, "c")
+            else:
+                self.fail("Unexpected x row in shuffled output.")
+
 
 class RepeatRandomSamplerTester(unittest.TestCase):
     def test_sampler(self):
@@ -216,6 +251,141 @@ def test_sampler_with_mini_repeat_count_and_batch_size_3(self):
         assert sampled[24:28] == sampled[28:32] == sampled[32:36]
 
 
+class TruncateWithProtectedTokensTester(unittest.TestCase):
+    def test_basic_example(self):
+        """Test the basic example from the problem description."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [2, 3, 6]
+        target_length = 3
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        expected_ids = torch.tensor([[2, 3, 5], [6, 9, 10]])
+        expected_mask = torch.ones_like(expected_ids)
+
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+        self.assertTrue(torch.equal(new_mask, expected_mask))
+
+    def test_no_truncation_needed(self):
+        """Test when target length equals current length."""
+        prompt_ids = torch.tensor([[1, 2, 3]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [2]
+        target_length = 3
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        self.assertTrue(torch.equal(new_ids, prompt_ids))
+        self.assertTrue(torch.equal(new_mask, prompt_mask))
+
+    def test_no_protected_tokens(self):
+        """Test truncation with no protected tokens (normal right truncation)."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = []
+        target_length = 3
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        expected_ids = torch.tensor([[3, 4, 5]])  # Last 3 tokens
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+
+    def test_all_tokens_protected(self):
+        """Test when all remaining tokens are protected."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [3, 4, 5]
+        target_length = 3
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        expected_ids = torch.tensor([[3, 4, 5]])
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+
+    def test_too_many_protected_tokens(self):
+        """Test error when too many protected tokens for target length."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [1, 2, 3, 4]
+        target_length = 3
+
+        with self.assertRaises(ValueError):
+            truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+    def test_single_batch_single_token(self):
+        """Test edge case with single batch and single token."""
+        prompt_ids = torch.tensor([[5]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [5]
+        target_length = 1
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        self.assertTrue(torch.equal(new_ids, prompt_ids))
+
+    def test_mask_preservation(self):
+        """Test that mask values are correctly preserved."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        prompt_mask = torch.tensor([[1, 0, 1, 0, 1]])  # Mixed mask values
+        protected_tokens = [2, 4]
+        target_length = 3
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        expected_ids = torch.tensor([[2, 4, 5]])
+        expected_mask = torch.tensor([[0, 0, 1]])  # Corresponding mask values
+
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+        self.assertTrue(torch.equal(new_mask, expected_mask))
+
+    def test_multiple_batches_different_protected(self):
+        """Test multiple batches where protected tokens appear differently."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5], [2, 6, 7, 8, 9], [10, 11, 12, 2, 13]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [2]
+        target_length = 3
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        expected_ids = torch.tensor(
+            [
+                [2, 4, 5],  # 2 is protected, keep last 2 non-protected (4,5)
+                [2, 8, 9],  # 2 is protected, keep last 2 non-protected (8,9)
+                [12, 2, 13],  # 2 is protected, keep last 2 non-protected (12,13)
+            ]
+        )
+
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+
+    def test_order_preservation(self):
+        """Test that relative order is preserved."""
+        prompt_ids = torch.tensor([[10, 2, 20, 3, 30, 40]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = [2, 3]
+        target_length = 4
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        # Should keep protected tokens 2,3 and last 2 non-protected tokens 30,40
+        # Order should be: 2, 3, 30, 40 (maintaining original relative positions)
+        expected_ids = torch.tensor([[2, 3, 30, 40]])
+
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+
+    def test_empty_protected_tokens_list(self):
+        """Test with empty protected tokens list."""
+        prompt_ids = torch.tensor([[1, 2, 3, 4, 5]])
+        prompt_mask = torch.ones_like(prompt_ids)
+        protected_tokens = []
+        target_length = 2
+
+        new_ids, new_mask = truncate_with_protected_tokens(prompt_ids, prompt_mask, target_length, protected_tokens)
+
+        expected_ids = torch.tensor([[4, 5]])  # Last 2 tokens
+        self.assertTrue(torch.equal(new_ids, expected_ids))
+
+
 class GetHighEntropyMaskTester(unittest.TestCase):
     def test_compute_entropy_mask_0(self):
         # We have a total of 12 tokens out of which 10 are non-pad.
@@ -270,6 +440,60 @@ def test_compute_entropy_all_masked(self):
         torch.testing.assert_close(entropy_mask, expected_mask)
 
 
+class SplitPixelValuesByGridTester(unittest.TestCase):
+    def test_split_correctly_0(self):
+        batch = {
+            "image_grid_thw": torch.tensor([[1, 2, 2], [1, 1, 4]]),  # Products: [4, 4]
+            "pixel_values": torch.arange(8 * 3).reshape(8, 3),  # Shape: [8, 3]
+        }
+        result = split_pixel_values_by_grid(batch)
+        self.assertIsInstance(result["pixel_values"], list)
+        self.assertEqual(len(result["pixel_values"]), 2)
+        self.assertTrue(torch.equal(result["pixel_values"][0], batch["pixel_values"][:4]))
+        self.assertTrue(torch.equal(result["pixel_values"][1], batch["pixel_values"][4:]))
+
+    def test_split_correctly_1(self):
+        batch = {
+            "image_grid_thw": torch.tensor([[1, 2, 2], [1, 2, 4]]),  # Products: [4, 8]
+            "pixel_values": torch.arange(12 * 3).reshape(12, 3),  # Shape: [12, 3]
+        }
+        result = split_pixel_values_by_grid(batch)
+        self.assertIsInstance(result["pixel_values"], list)
+        self.assertEqual(len(result["pixel_values"]), 2)
+        self.assertTrue(torch.equal(result["pixel_values"][0], batch["pixel_values"][:4]))
+        self.assertTrue(torch.equal(result["pixel_values"][1], batch["pixel_values"][4:12]))
+
+    def test_missing_keys(self):
+        batch = {"pixel_values": torch.tensor([1.0])}
+        result = split_pixel_values_by_grid(batch)
+        self.assertEqual(result, batch)
+
+    def test_mismatched_length(self):
+        batch = {
+            "image_grid_thw": torch.tensor([[2, 1, 1], [2, 1, 1]]),  # Total = 4
+            "pixel_values": torch.randn(3, 5),  # Only 3 rows
+        }
+        with self.assertRaises(ValueError):
+            split_pixel_values_by_grid(batch)
+
+
+class UnsplitPixelValuesByGridTester(unittest.TestCase):
+    def test_unsplit_correctly(self):
+        split = [torch.randn(4, 5), torch.randn(2, 5)]
+        merged = torch.cat(split, dim=0)
+        batch = {"pixel_values": split, "other_key": torch.tensor([1])}
+        result = unsplit_pixel_values_by_grid(batch)
+        self.assertIsInstance(result["pixel_values"], torch.Tensor)
+        self.assertTrue(torch.allclose(result["pixel_values"], merged))
+        self.assertIn("other_key", result)
+
+    def test_no_op_if_not_list(self):
+        original = torch.randn(5, 3)
+        batch = {"pixel_values": original}
+        result = unsplit_pixel_values_by_grid(batch)
+        self.assertTrue(torch.equal(result["pixel_values"], original))
+
+
 class GRPOTrainerTester(unittest.TestCase):
     def test_init_minimal(self):
         # Test that GRPOTrainer can be instantiated with only model, reward_model and train_dataset
@@ -800,40 +1024,6 @@ def reward_func(completions, some_values, **kwargs):
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
 
-    @require_vllm
-    @unittest.skip("We should add a mock for the vLLM server.")
-    def test_training_vllm(self):
-        """Test that training works with vLLM for generation."""
-        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
-
-        with tempfile.TemporaryDirectory() as tmp_dir:
-            training_args = GRPOConfig(
-                output_dir=tmp_dir,
-                learning_rate=0.1,  # increase the learning rate to speed up the test
-                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
-                num_generations=3,  # reduce the number of generations to reduce memory usage
-                max_completion_length=8,  # reduce the completion length to reduce memory usage
-                report_to="none",
-                use_vllm=True,
-            )
-            trainer = GRPOTrainer(
-                model="Qwen/Qwen2.5-0.5B-Instruct",  # tiny is too small for vLLM
-                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
-                args=training_args,
-                train_dataset=dataset,
-            )
-
-            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
-
-            trainer.train()
-
-            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
-
-            # Check that the params have changed
-            for n, param in previous_trainable_params.items():
-                new_param = trainer.model.get_parameter(n)
-                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
-
     def test_training_with_sync_ref_model(self):
         dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
 
@@ -948,7 +1138,7 @@ def test_training_vllm_and_peft(self):
             )
             lora_config = LoraConfig(
                 target_modules="all-linear",
-                # test with non-default modules as it add extra keys in state_dict tht we need to handle
+                # test with non-default modules as it adds extra keys in state_dict that we need to handle
                 modules_to_save=["embed_tokens", "lm_head"],
             )
             trainer = GRPOTrainer(
@@ -1122,7 +1312,7 @@ def test_training_with_mask_truncated_completions(self, mock_generate):
         # We mock the generate method because the model's random weights make it extremely unlikely to produce a
         # sequence containing the EOS token within the allowed max_completion_length. As a result, all tokens are
         # masked in the loss, the model doesn't update, and the final check (which verifies the update) fails.
-        def fake_generate(prompt_ids, **kwargs):
+        def fake_generate(input_ids, **kwargs):
             # pad_token_id = 151643; eos_token_id = 151645
             completions_ids = torch.tensor(
                 [
@@ -1130,9 +1320,9 @@ def fake_generate(prompt_ids, **kwargs):
                     [9, 10, 11, 151645, 151643, 151643, 151643, 151643],  # this one contains eos
                     [12, 13, 14, 15, 16, 17, 18, 151645],  # particular case, eos is generated just within the limit
                 ],
-                device=prompt_ids.device,
+                device=input_ids.device,
             )
-            return torch.cat([prompt_ids, completions_ids], dim=1)
+            return torch.cat([input_ids, completions_ids], dim=1)
 
         mock_generate.side_effect = fake_generate
 
@@ -1406,3 +1596,219 @@ def test_prepare_input_called_with_correct_data(self):
                     assert mock_prepare.call_args_list[i].args[1] == expected_first_generation_batch
                 for i in range(8, 16):
                     assert mock_prepare.call_args_list[i].args[1] == expected_second_generation_batch
+
+    @parameterized.expand(
+        [
+            ("trl-internal-testing/tiny-Gemma3ForConditionalGeneration",),
+            ("trl-internal-testing/tiny-Qwen2VLForConditionalGeneration",),
+            ("trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",),
+        ]
+    )
+    @require_vision
+    def test_training_vlm(self, model_id):
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model=model_id,
+                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check that the params have changed
+            for n, param in previous_trainable_params.items():
+                # Because of the way the tiny models are initialized, the gradient does not flow properly through the
+                # vision parts of the model, so we skip them. Ideally, we should fix the init of these models.
+                if n.startswith(("model.vision_tower.", "model.multi_modal_projector.")):
+                    continue
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+    @require_vision
+    def test_training_vlm_beta_non_zero(self):
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                beta=0.1,  # set beta to non-zero value to test the case where the reference model is used
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model="trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",
+                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check that the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+    @require_vision
+    @require_peft
+    def test_training_vlm_peft(self):
+        model = AutoModelForImageTextToText.from_pretrained(
+            "trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration"
+        )
+        base_param_names = [f"base_model.model.{n}" for n, _ in model.named_parameters()]
+        dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model=model,
+                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                args=training_args,
+                train_dataset=dataset,
+                peft_config=LoraConfig(target_modules=["q_proj", "v_proj"]),
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check that the peft params have changed and the base model params have not changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                if n in base_param_names:  # We expect the base model params to be the same
+                    self.assertTrue(torch.allclose(param, new_param), f"Parameter {n} has changed.")
+                elif "base_layer" not in n:  # We expect the peft params to be different (except for the base layer)
+                    self.assertFalse(torch.allclose(param, new_param), f"Parameter {n} has not changed.")
+
+    @require_vision
+    def test_training_vlm_and_importance_sampling(self):
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                steps_per_generation=2,  # increase the steps per generation to trigger IS
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model="trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",
+                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check that the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+    @require_vision
+    @require_liger_kernel
+    def test_training_vlm_and_liger(self):
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                use_liger_loss=True,  # Enable Liger loss
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model="trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",
+                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check that the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+    @require_vision
+    def test_training_vlm_and_prompt_truncation(self):
+        # If not handled properly, prompt truncation may truncate image token
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_prompt_only", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            training_args = GRPOConfig(
+                output_dir=tmp_dir,
+                learning_rate=0.1,  # increase the learning rate to speed up the test
+                per_device_train_batch_size=3,  # reduce the batch size to reduce memory usage
+                num_generations=3,  # reduce the number of generations to reduce memory usage
+                max_completion_length=8,  # reduce the completion length to reduce memory usage
+                max_prompt_length=18,
+                report_to="none",
+            )
+            trainer = GRPOTrainer(
+                model="trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",
+                reward_funcs="trl-internal-testing/tiny-Qwen2ForSequenceClassification-2.5",
+                args=training_args,
+                train_dataset=dataset,
+            )
+
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            trainer.train()
+
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check that the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.equal(param, new_param), f"Parameter {n} has not changed.")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/trl/extras/vllm_client.py b/trl/extras/vllm_client.py
index 71341a99542..2e6c1f86337 100644
--- a/trl/extras/vllm_client.py
+++ b/trl/extras/vllm_client.py
@@ -13,9 +13,11 @@
 # limitations under the License.
 
 import atexit
+import base64
 import logging
 import socket
 import time
+from io import BytesIO
 from typing import Optional
 from urllib.parse import urlparse
 
@@ -166,6 +168,7 @@ def check_server(self, total_timeout: float = 0.0, retry_interval: float = 2.0):
     def generate(
         self,
         prompts: list[str],
+        images: Optional[list] = None,
         n: int = 1,
         repetition_penalty: float = 1.0,
         temperature: float = 1.0,
@@ -182,6 +185,8 @@ def generate(
         Args:
             prompts (`list[str]`):
                 List of text prompts for which the model will generate completions.
+            images (`list[PIL.Image]` or `None`, *optional*, defaults to `None`):
+                List of PIL Images to send along with the prompts.
             n (`int`, *optional*, defaults to `1`):
                 Number of completions to generate for each prompt.
             repetition_penalty (`float`, *optional*, defaults to `1.0`):
@@ -208,10 +213,21 @@ def generate(
                 List of lists of token IDs representing the model-generated completions for each prompt.
         """
         url = f"{self.base_url}/generate/"
+
+        def pil_to_base64(image):
+            buffer = BytesIO()
+            image.save(buffer, format="PNG")
+            img_bytes = buffer.getvalue()
+            return base64.b64encode(img_bytes).decode("utf-8")
+
+        # Convert PIL images to base64 strings
+        images = [pil_to_base64(img) for img in images] if images else None
+
         response = self.session.post(
             url,
             json={
                 "prompts": prompts,
+                "images": images,
                 "n": n,
                 "repetition_penalty": repetition_penalty,
                 "temperature": temperature,
diff --git a/trl/scripts/vllm_serve.py b/trl/scripts/vllm_serve.py
index ffde86590c6..807aa5109cf 100644
--- a/trl/scripts/vllm_serve.py
+++ b/trl/scripts/vllm_serve.py
@@ -13,17 +13,20 @@
 # limitations under the License.
 
 import argparse
+import base64
 import logging
 import os
 from collections.abc import Sequence
 from contextlib import asynccontextmanager
 from dataclasses import dataclass, field
+from io import BytesIO
 from itertools import chain
 from multiprocessing import Pipe, Process
 from multiprocessing.connection import Connection
 from typing import Optional
 
 import torch
+from transformers import is_vision_available
 
 from trl import TrlParser
 from trl.import_utils import (
@@ -47,6 +50,10 @@
     import uvicorn
 
 
+if is_vision_available():
+    from PIL import Image
+
+
 if is_vllm_available():
     from vllm import LLM, SamplingParams
     from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator
@@ -426,6 +433,7 @@ async def get_world_size():
 
     class GenerateRequest(BaseModel):
         prompts: list[str]
+        images: Optional[list[str]] = None
         n: int = 1
         repetition_penalty: float = 1.0
         temperature: float = 1.0
@@ -447,6 +455,7 @@ async def generate(request: GenerateRequest):
         Args:
             request (`GenerateRequest`):
                 - `prompts` (list of `str`): A list of prompts (text strings) for the model to generate completions.
+                - `images` (list of `str`, *optional*, default to `None`): A list of base64 encoded images to process along with prompts.
                 - `n` (`int`, *optional*, defaults to `1`): Number of completions to generate for each prompt.
                 - `repetition_penalty` (`float`, *optional*, defaults to `1.0`): Repetition penalty to apply during generation.
                 - `temperature` (`float`, *optional*, defaults to `1.0`): Temperature for sampling. Higher values lead to more random outputs.
@@ -471,6 +480,14 @@ async def generate(request: GenerateRequest):
         {"completion_ids": [[101, 102, 103], [201, 202, 203]]}
         ```
         """
+        request.images = request.images or [None] * len(request.prompts)
+
+        prompts = []
+        for prompt, image in zip(request.prompts, request.images):
+            row = {"prompt": prompt}
+            if image is not None:
+                row["multi_modal_data"] = {"image": Image.open(BytesIO(base64.b64decode(image)))}
+            prompts.append(row)
 
         # Guided decoding, if enabled
         if request.guided_decoding_regex is not None:
@@ -492,7 +509,7 @@ async def generate(request: GenerateRequest):
         sampling_params = SamplingParams(**generation_kwargs)
 
         # Evenly distribute prompts across DP ranks
-        chunked_prompts = chunk_list(request.prompts, script_args.data_parallel_size)
+        chunked_prompts = chunk_list(prompts, script_args.data_parallel_size)
 
         # Send the prompts to each worker
         for connection, prompts in zip(connections, chunked_prompts):
diff --git a/trl/trainer/grpo_config.py b/trl/trainer/grpo_config.py
index 3e4adc1906e..04aa895eef0 100644
--- a/trl/trainer/grpo_config.py
+++ b/trl/trainer/grpo_config.py
@@ -582,7 +582,14 @@ def __post_init__(self):
                 "'generation_batch_size' and 'steps_per_generation' can not be both configured at the same time"
             )
 
-        # Check if the effective batch size can be divided by the number of generations
+        # The generation batch must contain full prompt groups (no partials), so it must be divisible by
+        # num_generations.
+        if self.generation_batch_size % self.num_generations != 0:
+            raise ValueError(
+                f"generation_batch_size ({self.generation_batch_size}) must be divisible by num_generations "
+                f"({self.num_generations})."
+            )
+
         if self.num_generations < 2:
             raise ValueError(
                 "GRPO requires at least 2 generations per prompt to calculate the advantages. You provided "
diff --git a/trl/trainer/grpo_trainer.py b/trl/trainer/grpo_trainer.py
index 0da501b6004..c2592c61e6a 100644
--- a/trl/trainer/grpo_trainer.py
+++ b/trl/trainer/grpo_trainer.py
@@ -12,12 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import copy
+import inspect
 import os
 import re
 import textwrap
 import warnings
 from collections import defaultdict, deque
-from collections.abc import Sized
+from collections.abc import Sequence, Sized
 from contextlib import nullcontext
 from functools import partial
 from pathlib import Path
@@ -34,12 +36,14 @@
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from torch.utils.data import DataLoader, Sampler
 from transformers import (
-    AutoModelForCausalLM,
+    AutoConfig,
     AutoModelForSequenceClassification,
+    AutoProcessor,
     AutoTokenizer,
     GenerationConfig,
     PreTrainedModel,
     PreTrainedTokenizerBase,
+    ProcessorMixin,
     Trainer,
     TrainerCallback,
     is_wandb_available,
@@ -231,28 +235,34 @@ def split_tensor_dict(
     ]
 
 
-def shuffle_tensor_dict(tensor_dict: dict[str, Optional[torch.Tensor]]) -> dict[str, Optional[torch.Tensor]]:
+def shuffle_sequence_dict(seq_dict: dict[str, Optional[Sequence]]) -> dict[str, Optional[Sequence]]:
     """
-    Shuffles a dictionary of tensors along the first dimension in unison.
+    Shuffles all sequence-like values in a dictionary along the first dimension in unison.
 
     Example:
     ```python
     >>> x = torch.arange(6).reshape(3, 2)
-    >>> y = torch.arange(3).reshape(3, 1)
-    >>> tensor_dict = {"x": x, "y": y}
-    >>> shuffle_tensor_dict(tensor_dict)
+    >>> y = ["a", "b", "c"]
+    >>> seq_dict = {"x": x, "y": y}
+    >>> shuffle_sequence_dict(seq_dict)
     {'x': tensor([[2, 3],
-                    [0, 1],
-                    [4, 5]]),
-        'y': tensor([[1],
-                    [0],
-                    [2]])}
+                  [0, 1],
+                  [4, 5]]),
+     'y': ['b', 'a', 'c']}
     ```
     """
-    first_tensor = next(tensor for tensor in tensor_dict.values() if tensor is not None)
-    batch_size = first_tensor.shape[0]
+    # Determine batch size from the first non-None sequence
+    batch_size = len(next(v for v in seq_dict.values() if v is not None))
     permutation = torch.randperm(batch_size)
-    return {key: tensor[permutation] if tensor is not None else None for key, tensor in tensor_dict.items()}
+
+    def permute(v: Optional[Sequence]) -> Optional[Sequence]:
+        if v is None:
+            return None
+        if isinstance(v, torch.Tensor):
+            return v[permutation]
+        return [v[i] for i in permutation]
+
+    return {key: permute(val) for key, val in seq_dict.items()}
 
 
 def nanmin(tensor: torch.Tensor) -> torch.Tensor:
@@ -290,6 +300,38 @@ def identity(x):
     return x
 
 
+def split_pixel_values_by_grid(batch: dict[str, torch.Tensor]) -> dict[str, Union[torch.Tensor, list[torch.Tensor]]]:
+    """
+    Splits `batch["pixel_values"]` into a list of tensors based on the product of each row in
+    `batch["image_grid_thw"]`, while keeping other entries unchanged.
+    """
+    if "image_grid_thw" not in batch or "pixel_values" not in batch:
+        return batch
+
+    lengths = batch["image_grid_thw"].prod(dim=1).tolist()  # [batch_size]
+    pixel_values = batch["pixel_values"]  # [total, feature_dim]
+
+    if sum(lengths) != pixel_values.size(0):
+        raise ValueError(f"Mismatch: sum(lengths) = {sum(lengths)} != pixel_values.size(0) = {pixel_values.size(0)}")
+
+    split_values = list(torch.split(batch["pixel_values"], lengths, dim=0))
+    return {**batch, "pixel_values": split_values}
+
+
+def unsplit_pixel_values_by_grid(batch: dict[str, Union[torch.Tensor, list[torch.Tensor]]]) -> dict[str, torch.Tensor]:
+    """
+    Opposite of `split_pixel_values_by_grid`. Merges a list of tensors in `batch["pixel_values"]`
+    back into a single tensor along the first dimension.
+    """
+    pixel_values = batch.get("pixel_values")
+
+    if isinstance(pixel_values, list):
+        merged = torch.cat(pixel_values, dim=0)
+        return {**batch, "pixel_values": merged}
+    else:
+        return batch
+
+
 def get_high_entropy_mask(entropies: torch.Tensor, mask: torch.Tensor, threshold: float) -> torch.Tensor:
     """
     Returns a binary mask identifying tokens whose entropy exceeds a given quantile threshold.
@@ -316,6 +358,63 @@ def get_high_entropy_mask(entropies: torch.Tensor, mask: torch.Tensor, threshold
     return entropy_mask & mask.bool()  # ensure padding tokens are always masked out
 
 
+def truncate_with_protected_tokens(
+    ids: torch.Tensor, mask: torch.Tensor, target_length: int, protected_tokens: list[int]
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Truncate tensors to target length while preserving protected tokens.
+
+    Args:
+        ids (`torch.Tensor`):
+            Input tensor of token IDs, shape (batch_size, sequence_length).
+        mask (`torch.Tensor`):
+            Input tensor of attention masks, shape (batch_size, sequence_length).
+        target_length (`int`):
+            Desired length of the output sequences.
+        protected_tokens (`list[int]`):
+            List of token IDs that should be preserved in the output.
+    """
+    protected_set = set(protected_tokens)
+
+    def process_sequence(ids, mask):
+        # Create boolean masks
+        is_protected = torch.tensor([x.item() in protected_set for x in ids])
+        is_non_protected = ~is_protected
+
+        # Count tokens
+        num_protected = is_protected.sum().item()
+        num_non_protected_needed = target_length - num_protected
+
+        if num_non_protected_needed < 0:
+            raise ValueError(
+                f"target_length ({target_length}) is too small for the protected tokens ({num_protected} tokens). "
+                f"Please increase target length to at least {num_protected} or disable truncation."
+            )
+
+        # Select which non-protected tokens to keep (rightmost ones)
+        non_protected_indices = torch.where(is_non_protected)[0]
+        keep_non_protected = torch.zeros_like(is_non_protected)
+        if num_non_protected_needed > 0:
+            keep_indices = non_protected_indices[-num_non_protected_needed:]
+            keep_non_protected[keep_indices] = True
+
+        # Final mask: protected OR selected non-protected
+        keep_mask = is_protected | keep_non_protected
+
+        return ids[keep_mask], mask[keep_mask]
+
+    # Process each sequence in the batch
+    truncated_seq = []
+    truncated_mask = []
+
+    for i in range(ids.shape[0]):
+        new_ids, new_mask = process_sequence(ids[i], mask[i])
+        truncated_seq.append(new_ids)
+        truncated_mask.append(new_mask)
+
+    return torch.stack(truncated_seq), torch.stack(truncated_mask)
+
+
 class GRPOTrainer(Trainer):
     """
     Trainer for the Group Relative Policy Optimization (GRPO) method. This algorithm was initially proposed in the
@@ -368,16 +467,15 @@ def reward_func(completions, **kwargs):
                 - A [`~transformers.PreTrainedModel`] object: Only sequence classification models are supported.
                 - A custom reward function: The function is provided with the prompts and the generated completions,
                   plus any additional columns in the dataset. It should return a list of rewards. Custom reward
-                  functions can also return None when the reward is not applicable to those samples. This is useful for
-                  multi-task training where different reward functions apply to different types of samples. When a
-                  reward function returns None for a sample, that reward function is excluded from the reward
+                  functions can also return `None` when the reward is not applicable to those samples. This is useful
+                  for multi-task training where different reward functions apply to different types of samples. When a
+                  reward function returns `None` for a sample, that reward function is excluded from the reward
                   calculation for that sample. For more details, see [Using a custom reward
                   function](#using-a-custom-reward-function).
 
-                  The trainer's state is also passed to the reward function. The trainer's state is an
-                  instance of [`transformers.TrainerState`](https://huggingface.co/docs/transformers/main/main_classes/callback#transformers.TrainerState)
-                  and can be accessed by accessing the `trainer_state` argument
-                  to the reward function's signature.
+                  The trainer's state is also passed to the reward function. The trainer's state is an instance of
+                  [`~transformers.TrainerState`] and can be accessed by accessing the `trainer_state` argument to the
+                  reward function's signature.
             - A list of reward functions, where each item can independently be any of the above types. Mixing different
             types within the list (e.g., a string model ID and a custom reward function) is allowed.
         args ([`GRPOConfig`], *optional*, defaults to `None`):
@@ -391,11 +489,11 @@ def reward_func(completions, **kwargs):
               and content).
         eval_dataset ([`~datasets.Dataset`], [`~datasets.IterableDataset`] or `dict[str, Union[Dataset, IterableDataset]]`):
             Dataset to use for evaluation. It must meet the same requirements as `train_dataset`.
-        processing_class ([`~transformers.PreTrainedTokenizerBase`], *optional*, defaults to `None`):
+        processing_class ([`~transformers.PreTrainedTokenizerBase`] or [`~transformers.ProcessorMixin`], *optional*, defaults to `None`):
             Processing class used to process the data. The padding side must be set to "left". If `None`, the
-            processing class is loaded from the model's name with [`~transformers.AutoTokenizer.from_pretrained`]. A
-            padding token, `processing_class.pad_token`, must be set. If the processing class has not set a padding
-            token, `processing_class.eos_token` will be used as the default.
+            processing class is loaded from the model's name with [`~transformers.AutoProcessor.from_pretrained`]. A
+            padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token,
+            `tokenizer.eos_token` will be used as the default.
         reward_processing_classes (`Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]`, *optional*, defaults to `None`):
             Processing classes corresponding to the reward functions specified in `reward_funcs`. Can be either:
 
@@ -428,7 +526,7 @@ def __init__(
         args: Optional[GRPOConfig] = None,
         train_dataset: Optional[Union[Dataset, IterableDataset]] = None,
         eval_dataset: Optional[Union[Dataset, IterableDataset, dict[str, Union[Dataset, IterableDataset]]]] = None,
-        processing_class: Optional[PreTrainedTokenizerBase] = None,
+        processing_class: Optional[Union[PreTrainedTokenizerBase, ProcessorMixin]] = None,
         reward_processing_classes: Optional[Union[PreTrainedTokenizerBase, list[PreTrainedTokenizerBase]]] = None,
         callbacks: Optional[list[TrainerCallback]] = None,
         optimizers: tuple[Optional[torch.optim.Optimizer], Optional[torch.optim.lr_scheduler.LambdaLR]] = (None, None),
@@ -457,10 +555,9 @@ def __init__(
                     f"a `torch.dtype` (e.g., 'float32'), but got {torch_dtype}."
                 )
             # Disable caching if gradient checkpointing is enabled (not supported)
-            model_init_kwargs["use_cache"] = (
-                False if args.gradient_checkpointing else model_init_kwargs.get("use_cache")
-            )
-            model = AutoModelForCausalLM.from_pretrained(model, **model_init_kwargs)
+            config = AutoConfig.from_pretrained(model_id)
+            architecture = getattr(transformers, config.architectures[0])
+            model = architecture.from_pretrained(model_id, **model_init_kwargs)
         else:
             model_id = model.config._name_or_path
             if args.model_init_kwargs is not None:
@@ -469,6 +566,14 @@ def __init__(
                     "This argument can only be used when the `model` argument is a string."
                 )
 
+        # Some models (SmolVLM/Idefics3) don't support `logits_to_keep` argument and error out if we pass it
+        # Inspect the forward method before we wrap the model with PEFT
+        self.model_kwarg_keys = (
+            inspect.signature(model.forward).parameters.keys()
+            if not hasattr(model, "get_base_model")
+            else inspect.signature(model.get_base_model().forward).parameters.keys()
+        )
+
         if peft_config is not None:
             if not is_peft_available():
                 raise ImportError("PEFT is required to use `peft_config`. Run `pip install peft`.")
@@ -480,9 +585,26 @@ def __init__(
 
         # Processing class
         if processing_class is None:
-            processing_class = AutoTokenizer.from_pretrained(model.config._name_or_path, padding_side="left")
-        if processing_class.pad_token is None:
-            processing_class.pad_token = processing_class.eos_token
+            processing_class = AutoProcessor.from_pretrained(model.config._name_or_path, padding_side="left")
+
+        # Handle pad token for processors or tokenizers
+        if isinstance(processing_class, ProcessorMixin):
+            tokenizer = processing_class.tokenizer
+        elif isinstance(processing_class, PreTrainedTokenizerBase):
+            tokenizer = processing_class
+        else:
+            raise TypeError("The `processing_class` must be either a `PreTrainedTokenizerBase` or a `ProcessorMixin`")
+
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
+
+        self.pad_token = tokenizer.pad_token
+        self.pad_token_id = tokenizer.pad_token_id
+        self.eos_token_id = tokenizer.eos_token_id
+        self.image_token = getattr(processing_class, "image_token", None)
+        self.image_token_id = getattr(processing_class, "image_token_id", None)
+        self.vision_start_token_id = getattr(model.config, "vision_start_token_id", None)
+        self.vision_end_token_id = getattr(model.config, "vision_end_token_id", None)
 
         # Reward functions
         if not isinstance(reward_funcs, list):
@@ -610,7 +732,9 @@ def __init__(
             self.ref_model = None
         else:
             # For deepspeed, fsdp or non-distributed models, create a reference model from scratch
-            self.ref_model = AutoModelForCausalLM.from_pretrained(model_id, **model_init_kwargs)
+            config = AutoConfig.from_pretrained(model_id)
+            architecture = getattr(transformers, config.architectures[0])
+            self.ref_model = architecture.from_pretrained(model_id, **model_init_kwargs)
 
         # Disable dropout in the models
         if args.disable_dropout:
@@ -644,7 +768,8 @@ def __init__(
         self.wandb_log_unique_prompts = args.wandb_log_unique_prompts
         self.num_completions_to_print = args.num_completions_to_print
         # Keep logs sized to the generation batch to record only outputs from the latest model update.
-        self._textual_logs = {
+        self._logs = {
+            "image": deque(maxlen=args.generation_batch_size),
             "prompt": deque(maxlen=args.generation_batch_size),
             "completion": deque(maxlen=args.generation_batch_size),
             "rewards": defaultdict(lambda: deque(maxlen=args.generation_batch_size)),
@@ -697,6 +822,10 @@ def __init__(
                 os.environ["MASTER_ADDR"] = os.environ.get("MASTER_ADDR", "localhost")
                 os.environ["MASTER_PORT"] = os.environ.get("MASTER_PORT", "12345")
 
+                if self.max_prompt_length is not None and self.max_completion_length is not None:
+                    max_model_len = self.max_prompt_length + self.max_completion_length
+                else:
+                    max_model_len = None
                 self.llm = LLM(
                     model=model.name_or_path,
                     tensor_parallel_size=args.vllm_tensor_parallel_size,
@@ -704,7 +833,7 @@ def __init__(
                     max_num_seqs=self.args.per_device_train_batch_size
                     * self.vllm_tensor_parallel_size
                     * self.args.steps_per_generation,
-                    max_model_len=self.max_prompt_length + self.max_completion_length,
+                    max_model_len=max_model_len,
                     distributed_executor_backend="external_launcher",
                     # Feed identical seed for tp groups to ensure sampling results are the same across workers
                     seed=self.accelerator.process_index // self.vllm_tensor_parallel_size,
@@ -725,9 +854,9 @@ def __init__(
             generation_kwargs = {
                 "max_new_tokens": self.max_completion_length,
                 "do_sample": True,
-                "pad_token_id": processing_class.pad_token_id,
-                "bos_token_id": processing_class.bos_token_id,
-                "eos_token_id": processing_class.eos_token_id,
+                "pad_token_id": tokenizer.pad_token_id,
+                "bos_token_id": tokenizer.bos_token_id,
+                "eos_token_id": tokenizer.eos_token_id,
                 "temperature": self.temperature,
                 "top_p": self.top_p,
                 "top_k": self.top_k,
@@ -778,7 +907,7 @@ def _set_signature_columns_if_needed(self):
         # In GRPOTrainer, we preprocess data, so using the model's signature columns doesn't work.
         # Instead, we set them to the columns expected by the `training_step` method, hence the override.
         if self._signature_columns is None:
-            self._signature_columns = ["prompt"]
+            self._signature_columns = ["prompt", "image"]
 
     # This method overrides `Trainer.get_train_dataloader` to support our custom batching strategy.
     # Instead of returning a standard per-step batch (i.e., `per_device_batch_size), our dataloader loads an
@@ -890,18 +1019,43 @@ def _enable_gradient_checkpointing(self, model: PreTrainedModel, args: GRPOConfi
         return model
 
     @profiling_decorator
-    def _get_last_hidden_state(self, unwrapped_model, input_ids, attention_mask, logits_to_keep=None):
+    def _get_last_hidden_state(
+        self, unwrapped_model, input_ids, attention_mask, logits_to_keep, pixel_values, image_grid_thw
+    ):
         if is_peft_model(unwrapped_model):
             unwrapped_model = unwrapped_model.base_model.model
-        last_hidden_state = unwrapped_model.model(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+
+        # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't)
+        model_inputs = {"input_ids": input_ids, "attention_mask": attention_mask}
+
+        if image_grid_thw is not None and pixel_values is not None:
+            model_inputs["image_grid_thw"] = image_grid_thw
+        if pixel_values is not None:
+            model_inputs["pixel_values"] = pixel_values
+
+        # Only add logits_to_keep if the model supports it
+        if "logits_to_keep" in self.model_kwarg_keys:
+            # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
+            model_inputs["logits_to_keep"] = logits_to_keep + 1
+
+        last_hidden_state = unwrapped_model.model(**model_inputs).last_hidden_state
+        # Exclude the last value: it corresponds to the next token pred
         last_hidden_state = last_hidden_state[:, :-1, :]  # (B, L-1, H)
-        if logits_to_keep is not None:
-            last_hidden_state = last_hidden_state[:, -logits_to_keep:, :]  # (B, logits_to_keep, H)
+        # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op.
+        last_hidden_state = last_hidden_state[:, -logits_to_keep:, :]  # (B, logits_to_keep, H)
         return last_hidden_state
 
     @profiling_decorator
     def _get_per_token_logps_and_entropies(
-        self, model, input_ids, attention_mask, logits_to_keep, batch_size=None, compute_entropy=False
+        self,
+        model,
+        input_ids,
+        attention_mask,
+        logits_to_keep,
+        batch_size=None,
+        compute_entropy=False,
+        pixel_values=None,
+        image_grid_thw=None,
     ) -> dict[str, Optional[torch.Tensor]]:
         """Compute log‐probs and (optionally) entropies for each token."""
         batch_size = batch_size or input_ids.size(0)  # Chunk inputs into smaller batches to reduce memory peak
@@ -911,13 +1065,27 @@ def _get_per_token_logps_and_entropies(
             input_ids_batch = input_ids[start : start + batch_size]
             attention_mask_batch = attention_mask[start : start + batch_size]
 
-            # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
-            logits = model(
-                input_ids=input_ids_batch,
-                attention_mask=attention_mask_batch,
-                logits_to_keep=logits_to_keep + 1,
-            ).logits
-            logits = logits[:, :-1, :]  # (B, L-1, V), exclude the last logit: it corresponds to the next token pred
+            # Build model inputs - check if the model supports logits_to_keep (some models and VLMs don't)
+            model_inputs = {"input_ids": input_ids_batch, "attention_mask": attention_mask_batch}
+
+            if image_grid_thw is not None and pixel_values is not None:
+                model_inputs["image_grid_thw"] = image_grid_thw[start : start + batch_size]
+                start_pixel_idx = image_grid_thw[:start].prod(-1).sum().item()
+                end_pixel_idx = image_grid_thw[: start + batch_size].prod(-1).sum().item()
+                model_inputs["pixel_values"] = pixel_values[start_pixel_idx:end_pixel_idx]
+            elif pixel_values is not None:
+                model_inputs["pixel_values"] = pixel_values[start : start + batch_size]
+
+            # Only add logits_to_keep if the model supports it
+            if "logits_to_keep" in self.model_kwarg_keys:
+                # We add 1 to `logits_to_keep` because the last logits of the sequence is later excluded
+                model_inputs["logits_to_keep"] = logits_to_keep + 1
+
+            logits = model(**model_inputs).logits
+            # Exclude the last value: it corresponds to the next token pred
+            logits = logits[:, :-1, :]  # (B, L-1, H)
+            # Only keep the last logits_to_keep. For model that support logits_to_keep, this is a no-op.
+            logits = logits[:, -logits_to_keep:, :]  # (B, logits_to_keep, H)
             # Divide logits by sampling temperature.
             # See https://huggingface.co/blog/the_n_implementation_details_of_rlhf_with_ppo#policy-training-implementation-details
             logits = logits / self.temperature
@@ -1055,8 +1223,10 @@ def _prepare_inputs(
             if self._step % generate_every == 0 or self._buffered_inputs is None:
                 # self._buffered_inputs=None can occur when resuming from a checkpoint
                 generation_batch = self._generate_and_score_completions(generation_batch)
-                generation_batch = shuffle_tensor_dict(generation_batch)
-                self._buffered_inputs = split_tensor_dict(generation_batch, self.args.steps_per_generation)
+                generation_batch = split_pixel_values_by_grid(generation_batch)
+                generation_batch = shuffle_sequence_dict(generation_batch)
+                generation_batches = split_tensor_dict(generation_batch, self.args.steps_per_generation)
+                self._buffered_inputs = [unsplit_pixel_values_by_grid(batch) for batch in generation_batches]
             inputs = self._buffered_inputs[self._step % self.args.steps_per_generation]
             self._step += 1
         else:
@@ -1125,10 +1295,37 @@ def _generate_and_score_completions(
         mode = "train" if self.model.training else "eval"
 
         prompts = [x["prompt"] for x in inputs]
+
+        # We don't yet support visual reward models/function, so we keep a copy of the original text-only prompts for
+        # later use in the reward computation. If images are present, we insert {"type": "image"} as required by the
+        # VLM chat template.
+        original_prompts = copy.deepcopy(prompts)
+
+        # If the prompts are conversational and the inputs contain images, we need to convert the prompts from
+        # [{"role": "user", "content": "What color is the sky?"}] to
+        # [{"role": "user", "content": [{"type": "image"}, {"type": "text", "text": "What color is the sky?"}]}]
+        kwargs = {}
+        has_images = "image" in inputs[0]
+        if has_images:
+            images = [example.get("image") for example in inputs]
+            kwargs = {"images": [[img] for img in images]}
+            for prompt in prompts:
+                if isinstance(prompt, list):
+                    for message in prompt:
+                        if isinstance(message, dict) and message.get("role") == "user":
+                            if isinstance(message.get("content"), str):
+                                message["content"] = [{"type": "image"}, {"type": "text", "text": message["content"]}]
+                            break
+
         prompts_text = [maybe_apply_chat_template(example, self.processing_class)["prompt"] for example in inputs]
 
         prompt_inputs = self.processing_class(
-            text=prompts_text, return_tensors="pt", padding=True, padding_side="left", add_special_tokens=False
+            text=prompts_text,
+            return_tensors="pt",
+            padding=True,
+            padding_side="left",
+            add_special_tokens=False,
+            **kwargs,
         )
         prompt_inputs = super()._prepare_inputs(prompt_inputs)
         prompt_ids, prompt_mask = prompt_inputs["input_ids"], prompt_inputs["attention_mask"]
@@ -1137,14 +1334,25 @@ def _generate_and_score_completions(
             # If max_prompt_length is set, we trim the prompt to keep only the last `max_prompt_length` tokens.
             # Then we decode those tokens back into text. We manually remove leading pad tokens from the decoded text,
             # because we can't use `skip_special_tokens=True` (some special tokens are still needed for generation).
-            prompt_ids = prompt_ids[:, -self.max_prompt_length :]
-            prompt_mask = prompt_mask[:, -self.max_prompt_length :]
+            protected = [self.image_token_id, self.vision_start_token_id, self.vision_end_token_id]
+            protected = [token for token in protected if token is not None]
+            prompt_ids, prompt_mask = truncate_with_protected_tokens(
+                prompt_ids, prompt_mask, self.max_prompt_length, protected
+            )
+
             prompts_text = self.processing_class.batch_decode(
                 prompt_ids, skip_special_tokens=False, clean_up_tokenization_spaces=False
             )
-            prompts_text = [
-                re.sub(rf"^({re.escape(self.processing_class.pad_token)})+", "", text) for text in prompts_text
-            ]
+            prompts_text = [re.sub(rf"^({re.escape(self.pad_token)})+", "", text) for text in prompts_text]
+
+            # The chat template inserts a single image token into the prompt text. However, when this text is later
+            # tokenized, the single image token string is expanded into multiple image token IDs, depending on the
+            # image size. Since we're detokenizing here, we may see repeated image tokens in the decoded text. We
+            # collapse them back into a single token string to match the original template.
+            if self.image_token is not None:
+                prompts_text = [
+                    re.sub(rf"({re.escape(self.image_token)})+", self.image_token, text) for text in prompts_text
+                ]
 
         # Generate completions using either vLLM or regular generation
         if self.use_vllm:
@@ -1156,14 +1364,24 @@ def _generate_and_score_completions(
             # Generate completions using vLLM: gather all prompts and use them in a single call in the main process
             if self.vllm_mode == "server":
                 all_prompts_text = gather_object(prompts_text)
+                if has_images:
+                    all_images = gather_object(images)
+
                 if self.accelerator.is_main_process:
                     # Since 'prompts' contains 'num_generations' duplicates, we first take unique prompts, and generate
                     # num_generations outputs for each one. This is faster than generating outputs for each duplicate
                     # prompt individually.
                     ordered_set_of_prompts = all_prompts_text[:: self.num_generations]
+
+                    if has_images:
+                        ordered_set_of_images = all_images[:: self.num_generations]
+                    else:
+                        ordered_set_of_images = None
+
                     with profiling_context(self, "vLLM.generate"):
                         completion_ids = self.vllm_client.generate(
                             prompts=ordered_set_of_prompts,
+                            images=ordered_set_of_images,
                             n=self.num_generations,
                             repetition_penalty=self.repetition_penalty,
                             temperature=self.temperature,
@@ -1188,7 +1406,7 @@ def _generate_and_score_completions(
             # Generate completions using colocated vLLM instances: each device holds vLLM copy and work on their own batch of prompts
             elif self.vllm_mode == "colocate":
                 if self.guided_decoding_regex:
-                    guided_decoding = GuidedDecodingParams(backend="outlines", regex=self.guided_decoding_regex)
+                    guided_decoding = GuidedDecodingParams(regex=self.guided_decoding_regex)
                 else:
                     guided_decoding = None
 
@@ -1213,11 +1431,29 @@ def _generate_and_score_completions(
                     gathered_prompts = [None for _ in range(self.vllm_tensor_parallel_size)]
                     torch.distributed.all_gather_object(gathered_prompts, prompts_text, group=self.tp_group)
                     all_prompts_text = [p for sublist in gathered_prompts for p in sublist]
+
+                    if has_images:
+                        gathered_images = [None for _ in range(self.vllm_tensor_parallel_size)]
+                        torch.distributed.all_gather_object(gathered_images, images, group=self.tp_group)
+                        all_images = [img for sublist in gathered_images for img in sublist]
+                    else:
+                        all_images = None
                 else:
                     all_prompts_text = prompts_text
+                    all_images = images if has_images else None
+
+                if has_images and all_images:
+                    vllm_inputs = []
+                    for prompt, image in zip(all_prompts_text, all_images):
+                        if image is not None:
+                            vllm_inputs.append({"prompt": prompt, "multi_modal_data": {"image": image}})
+                        else:
+                            vllm_inputs.append(prompt)
+                else:
+                    vllm_inputs = all_prompts_text
 
                 with profiling_context(self, "vLLM.generate"):
-                    all_outputs = self.llm.generate(all_prompts_text, sampling_params=sampling_params, use_tqdm=False)
+                    all_outputs = self.llm.generate(vllm_inputs, sampling_params=sampling_params, use_tqdm=False)
 
                 completion_ids = [output.token_ids for outputs in all_outputs for output in outputs.outputs]
 
@@ -1230,11 +1466,13 @@ def _generate_and_score_completions(
 
             # Pad the completions, and concatenate them with the prompts
             completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]
-            completion_ids = pad(completion_ids, padding_value=self.processing_class.pad_token_id)
+            completion_ids = pad(completion_ids, padding_value=self.pad_token_id)
             prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
 
         elif self.use_transformers_paged:
-            prompt_inputs = self.processing_class(text=prompts_text)
+            # Re-process inputs for paged generation if needed
+            # Note: images are already validated and preprocessed above
+            paged_prompt_inputs = self.processing_class(text=prompts_text, **kwargs)
             previous_attn = self.model_wrapped.config._attn_implementation
 
             if is_flash_attn_2_available():
@@ -1256,15 +1494,13 @@ def _generate_and_score_completions(
                     unwrapped_model.to(torch.float16)
                 with torch.inference_mode():
                     all_outputs = unwrapped_model.generate_batch(
-                        prompt_inputs.input_ids, generation_config=self.generation_config, progress_bar=False
+                        paged_prompt_inputs.input_ids, generation_config=self.generation_config, progress_bar=False
                     )
             completion_ids = [output.generated_tokens for output in all_outputs.values()]
             completion_ids = [torch.tensor(ids, device=device) for ids in completion_ids]
-            completion_ids = pad(
-                completion_ids, padding_value=self.processing_class.pad_token_id, padding_side="right"
-            )
-            prompt_ids = [torch.tensor(ids, device=device) for ids in prompt_inputs.input_ids]
-            prompt_ids = pad(prompt_ids, padding_value=self.processing_class.pad_token_id, padding_side="left")
+            completion_ids = pad(completion_ids, padding_value=self.pad_token_id, padding_side="right")
+            prompt_ids = [torch.tensor(ids, device=device) for ids in paged_prompt_inputs.input_ids]
+            prompt_ids = pad(prompt_ids, padding_value=self.pad_token_id, padding_side="left")
             prompt_completion_ids = torch.cat([prompt_ids, completion_ids], dim=1)
             # Restore the original attention implementation, training mode
             self.model_wrapped.config._attn_implementation = previous_attn
@@ -1278,17 +1514,17 @@ def _generate_and_score_completions(
                 torch.no_grad(),
                 FSDP.summon_full_params(self.model_wrapped, recurse=False) if self.is_fsdp_enabled else nullcontext(),
             ):
+                prompt_inputs["input_ids"], prompt_inputs["attention_mask"] = prompt_ids, prompt_mask
                 prompt_completion_ids = unwrapped_model.generate(
-                    prompt_ids, attention_mask=prompt_mask, generation_config=self.generation_config
+                    **prompt_inputs, generation_config=self.generation_config
                 )
-
             # Compute prompt length and extract completion ids
             prompt_length = prompt_ids.size(1)
             prompt_ids = prompt_completion_ids[:, :prompt_length]
             completion_ids = prompt_completion_ids[:, prompt_length:]
 
         # Mask everything after the first EOS token
-        is_eos = completion_ids == self.processing_class.eos_token_id
+        is_eos = completion_ids == self.eos_token_id
         eos_idx = torch.full((is_eos.size(0),), is_eos.size(1), dtype=torch.long, device=device)
         eos_idx[is_eos.any(dim=1)] = is_eos.int().argmax(dim=1)[is_eos.any(dim=1)]
         sequence_indices = torch.arange(is_eos.size(1), device=device).expand(is_eos.size(0), -1)
@@ -1323,7 +1559,13 @@ def _generate_and_score_completions(
             generate_every = self.args.steps_per_generation * self.num_iterations  # generation frequency
             if self.args.gradient_accumulation_steps % generate_every != 0:
                 old_per_token_logps, _ = self._get_per_token_logps_and_entropies(
-                    self.model, prompt_completion_ids, attention_mask, logits_to_keep, batch_size
+                    self.model,
+                    prompt_completion_ids,
+                    attention_mask,
+                    logits_to_keep,
+                    batch_size,
+                    pixel_values=prompt_inputs.get("pixel_values"),
+                    image_grid_thw=prompt_inputs.get("image_grid_thw"),
                 )
             else:
                 old_per_token_logps = None
@@ -1332,12 +1574,24 @@ def _generate_and_score_completions(
             if self.beta != 0.0:
                 if self.ref_model is not None:
                     ref_per_token_logps, _ = self._get_per_token_logps_and_entropies(
-                        self.ref_model, prompt_completion_ids, attention_mask, logits_to_keep, batch_size
+                        self.ref_model,
+                        prompt_completion_ids,
+                        attention_mask,
+                        logits_to_keep,
+                        batch_size=batch_size,
+                        pixel_values=prompt_inputs.get("pixel_values"),
+                        image_grid_thw=prompt_inputs.get("image_grid_thw"),
                     )
                 else:
                     with self.accelerator.unwrap_model(self.model).disable_adapter():
                         ref_per_token_logps, _ = self._get_per_token_logps_and_entropies(
-                            self.model, prompt_completion_ids, attention_mask, logits_to_keep, batch_size
+                            self.model,
+                            prompt_completion_ids,
+                            attention_mask,
+                            logits_to_keep,
+                            batch_size=batch_size,
+                            pixel_values=prompt_inputs.get("pixel_values"),
+                            image_grid_thw=prompt_inputs.get("image_grid_thw"),
                         )
             else:
                 ref_per_token_logps = None
@@ -1355,7 +1609,7 @@ def _generate_and_score_completions(
         # Calculate rewards for each reward function. rewards_per_func aggregates rewards across all processes. This is
         # important because rewards will be normalized per group, and completions are distributed. We will later slice
         # rewards_per_func to extract each process's subset.
-        rewards_per_func = self._calculate_rewards(inputs, prompts, completions, completion_ids_list)
+        rewards_per_func = self._calculate_rewards(inputs, original_prompts, completions, completion_ids_list)
 
         # Apply weights to each reward function's output and sum
         rewards = (rewards_per_func * self.reward_weights.to(device).unsqueeze(0)).nansum(dim=1)
@@ -1413,21 +1667,31 @@ def _generate_and_score_completions(
         self._metrics[mode]["frac_reward_zero_std"].append(is_std_zero.float().mean().item())
 
         # Log prompt and completion texts
-        self._textual_logs["prompt"].extend(gather_object(prompts_text))
-        self._textual_logs["completion"].extend(gather_object(completions_text))
+        self._logs["prompt"].extend(gather_object(prompts_text))
+        self._logs["completion"].extend(gather_object(completions_text))
         for i, name in enumerate(self.reward_func_names):
-            self._textual_logs["rewards"][name].extend(rewards_per_func[:, i].tolist())
-        self._textual_logs["advantages"].extend(all_process_advantages.tolist())
+            self._logs["rewards"][name].extend(rewards_per_func[:, i].tolist())
+        self._logs["advantages"].extend(all_process_advantages.tolist())
+
+        if has_images:
+            self._logs["image"].extend(gather_object(images))
 
-        return {
+        output = {
             "prompt_ids": prompt_ids,
             "prompt_mask": prompt_mask,
             "completion_ids": completion_ids,
             "completion_mask": completion_mask,
             "advantages": advantages,
-            "old_per_token_logps": old_per_token_logps,
-            "ref_per_token_logps": ref_per_token_logps,
         }
+        if old_per_token_logps is not None:
+            output["old_per_token_logps"] = old_per_token_logps
+        if ref_per_token_logps is not None:
+            output["ref_per_token_logps"] = ref_per_token_logps
+        if "pixel_values" in prompt_inputs:
+            output["pixel_values"] = prompt_inputs["pixel_values"]
+        if "image_grid_thw" in prompt_inputs:
+            output["image_grid_thw"] = prompt_inputs["image_grid_thw"]
+        return output
 
     def compute_liger_loss(self, unwrapped_model, inputs):
         # Compute the per-token log probabilities for the model
@@ -1437,8 +1701,15 @@ def compute_liger_loss(self, unwrapped_model, inputs):
         attention_mask = torch.cat([prompt_mask, completion_mask], dim=1)
         logits_to_keep = completion_ids.size(1)  # we only need to compute the logits for the completion tokens
 
-        # get the last hidden state of the model
-        last_hidden_state = self._get_last_hidden_state(unwrapped_model, input_ids, attention_mask, logits_to_keep)
+        # Get the last hidden state of the model
+        last_hidden_state = self._get_last_hidden_state(
+            unwrapped_model,
+            input_ids,
+            attention_mask,
+            logits_to_keep,
+            inputs.get("pixel_values"),
+            inputs.get("image_grid_thw"),
+        )
 
         # compute loss and metrics using liger grpo loss
         loss, metrics = self.liger_grpo_loss(
@@ -1448,8 +1719,8 @@ def compute_liger_loss(self, unwrapped_model, inputs):
             attention_mask=completion_mask,
             advantages=inputs["advantages"],
             bias=unwrapped_model.lm_head.bias,
-            old_per_token_logps=inputs["old_per_token_logps"],
-            ref_per_token_logps=inputs["ref_per_token_logps"],
+            old_per_token_logps=inputs.get("old_per_token_logps"),
+            ref_per_token_logps=inputs.get("ref_per_token_logps"),
         )
         # Extract metrics from the liger_grpo_loss output
         # KL divergence is the first metric when beta is non-zero
@@ -1483,7 +1754,13 @@ def _compute_loss(self, model, inputs):
 
         # Compute the per_token_logps and the entropy (if necessary) at each position in the completion
         per_token_logps, entropies = self._get_per_token_logps_and_entropies(
-            model, input_ids, attention_mask, logits_to_keep, compute_entropy=self.top_entropy_quantile < 1.0
+            model,
+            input_ids,
+            attention_mask,
+            logits_to_keep,
+            compute_entropy=self.top_entropy_quantile < 1.0,
+            pixel_values=inputs.get("pixel_values"),
+            image_grid_thw=inputs.get("image_grid_thw"),
         )
 
         if self.top_entropy_quantile < 1.0:
@@ -1503,9 +1780,8 @@ def _compute_loss(self, model, inputs):
         # When using num_iterations == 1 and steps_per_generation <= gradient_accumulation_steps
         # old_per_token_logps == per_token_logps, so we can skip it's computation
         # (see _generate_and_score_completions) and use per_token_logps.detach() instead.
-        old_per_token_logps = (
-            per_token_logps.detach() if inputs["old_per_token_logps"] is None else inputs["old_per_token_logps"]
-        )
+        old_per_token_logps = inputs.get("old_per_token_logps")
+        old_per_token_logps = per_token_logps.detach() if old_per_token_logps is None else old_per_token_logps
         coef_1 = torch.exp(per_token_logps - old_per_token_logps)
         coef_2 = torch.clamp(coef_1, 1 - self.epsilon_low, 1 + self.epsilon_high)
 
@@ -1580,10 +1856,10 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non
         if self.accelerator.is_main_process and self.log_completions:
             if is_rich_available():
                 print_prompt_completions_sample(
-                    self._textual_logs["prompt"],
-                    self._textual_logs["completion"],
-                    self._textual_logs["rewards"],
-                    self._textual_logs["advantages"],
+                    self._logs["prompt"],
+                    self._logs["completion"],
+                    self._logs["rewards"],
+                    self._logs["advantages"],
                     self.state.global_step,
                     self.num_completions_to_print,
                 )
@@ -1592,12 +1868,22 @@ def log(self, logs: dict[str, float], start_time: Optional[float] = None) -> Non
                 import pandas as pd
 
                 table = {
-                    "step": [str(self.state.global_step)] * len(self._textual_logs["prompt"]),
-                    "prompt": self._textual_logs["prompt"],
-                    "completion": self._textual_logs["completion"],
-                    **self._textual_logs["rewards"],
-                    "advantage": self._textual_logs["advantages"],
+                    "step": [str(self.state.global_step)] * len(self._logs["prompt"]),
+                    "prompt": self._logs["prompt"],
+                    "completion": self._logs["completion"],
+                    **self._logs["rewards"],
+                    "advantage": self._logs["advantages"],
                 }
+
+                if self._logs["image"]:
+                    table["image"] = []
+                    for img in self._logs["image"]:
+                        if img is not None:
+                            # Convert images to wandb Image objects for proper visualization
+                            table["image"].append(wandb.Image(img))
+                        else:
+                            table["image"].append(None)
+
                 df = pd.DataFrame(table)
                 if self.wandb_log_unique_prompts:
                     df = df.drop_duplicates(subset=["prompt"])