🔮 Native VLM support for SFTTrainer

examples/scripts/sft_video_llm.py

@@ -179,7 +179,6 @@ class CustomScriptArguments(ScriptArguments):
     # Configure training args
     training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)
     training_args.remove_unused_columns = False
-    training_args.dataset_kwargs = {"skip_prepare_dataset": True}
 
     # Load dataset
     dataset = load_dataset(script_args.dataset_name, name=script_args.dataset_config, split="train")

examples/scripts/sft_vlm.py

@@ -31,7 +31,7 @@
     --per_device_train_batch_size 8 \
     --gradient_accumulation_steps 8 \
     --output_dir sft-llava-1.5-7b-hf \
-    --bf16 \
+    --bf16 True \
     --torch_dtype bfloat16 \
     --gradient_checkpointing
 
@@ -63,7 +63,6 @@
     script_args, training_args, model_args = parser.parse_args_and_config()
     training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)
     training_args.remove_unused_columns = False
-    training_args.dataset_kwargs = {"skip_prepare_dataset": True}
 
     ################
     # Model, Tokenizer & Processor

examples/scripts/sft_vlm_gemma3.py

@@ -30,7 +30,7 @@
     --per_device_train_batch_size 1 \
     --gradient_accumulation_steps 1 \
     --output_dir gemma-3-4b-it-trl-sft-llava-instruct-mix-vsft \
-    --bf16 \
+    --bf16 True \
     --torch_dtype bfloat16 \
     --use_peft \
     --lora_target_modules all-linear \
@@ -47,7 +47,7 @@
     --per_device_train_batch_size 1 \
     --gradient_accumulation_steps 1 \
     --output_dir gemma-3-4b-it-trl-sft-MMIU-Benchmark \
-    --bf16 \
+    --bf16 True \
     --torch_dtype bfloat16 \
     --use_peft \
     --lora_target_modules all-linear
@@ -142,7 +142,6 @@ def main():
     script_args, training_args, model_args = parser.parse_args_and_config()
     training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)
     training_args.remove_unused_columns = False
-    training_args.dataset_kwargs = {"skip_prepare_dataset": True}
 
     ################
     # Model, Tokenizer & Processor

examples/scripts/sft_vlm_smol_vlm.py

@@ -31,7 +31,7 @@
     --per_device_train_batch_size 1 \
     --gradient_accumulation_steps 1 \
     --output_dir sft-smol-vlm-hf \
-    --bf16 \
+    --bf16 True \
     --torch_dtype bfloat16 \
     --gradient_checkpointing \
     --use_peft \
@@ -70,7 +70,6 @@
     script_args, training_args, model_args = parser.parse_args_and_config()
     training_args.gradient_checkpointing_kwargs = dict(use_reentrant=False)
     training_args.remove_unused_columns = False
-    training_args.dataset_kwargs = {"skip_prepare_dataset": True}
 
     ################
     # Model, Tokenizer & Processor

scripts/generate_tiny_models.py

@@ -98,7 +98,7 @@
 api = HfApi()
 
 
-def push_to_hub(model, tokenizer, prefix=None, suffix=None):
+def push_to_hub(model, tokenizer, prefix=None, suffix=None, force=False):
     model_class_name = model.__class__.__name__
     content = MODEL_CARD.format(model_class_name=model_class_name)
     model_card = ModelCard(content)
@@ -108,7 +108,7 @@ def push_to_hub(model, tokenizer, prefix=None, suffix=None):
     if suffix is not None:
         repo_id += f"-{suffix}"
 
-    if api.repo_exists(repo_id):
+    if api.repo_exists(repo_id) and not force:
         print(f"Model {repo_id} already exists, skipping")
     else:
         model.push_to_hub(repo_id)
@@ -297,19 +297,26 @@ def init_weights_tiny_model(model):
         kwargs["projection_dim"] = 8
     if config_class in [LlavaConfig, LlavaNextConfig, PaliGemmaConfig]:
         vision_kwargs["projection_dim"] = 8
-    if config_class in [LlavaConfig, LlavaNextConfig]:
+    if config_class in [LlavaConfig, LlavaNextConfig, Gemma3Config]:
         vision_kwargs["image_size"] = 336
-        vision_kwargs["patch_size"] = 14
+        vision_kwargs["patch_size"] = 20
+        processor.image_processor.size = {"height": 336, "width": 336}
     if config_class in [Qwen2VLConfig, Qwen2_5_VLConfig]:
         kwargs["vision_start_token_id"] = 151652
-        text_kwargs["rope_scaling"] = {"type": "mrope", "mrope_section": [1]}
+        kwargs["vision_end_token_id"] = 151653
+        kwargs["vision_token_id"] = 151654
+        kwargs["image_token_id"] = 151655
+        kwargs["vocab_size"] = len(processor.tokenizer.vocab)
+        text_kwargs["rope_scaling"] = {"type": "mrope", "mrope_section": [2]}
         vision_kwargs["depth"] = 4
         vision_kwargs["embed_dim"] = 64
+    if config_class in [Qwen2_5_VLConfig]:
+        vision_kwargs["out_hidden_size"] = 16
 
     config = config_class(
         text_config=dict(
             vocab_size=processor.tokenizer.vocab_size + len(processor.tokenizer.added_tokens_encoder),
-            hidden_size=8,
+            hidden_size=16,
             num_attention_heads=4,
             num_key_value_heads=2,
             num_hidden_layers=2,
@@ -326,4 +333,4 @@ def init_weights_tiny_model(model):
         **kwargs,
     )
     model = model_class(config)
-    push_to_hub(model, processor, "tiny")
+    push_to_hub(model, processor, "tiny", force=True)

tests/test_sft_trainer.py

@@ -1472,3 +1472,31 @@ def test_train_with_torch_dtype(self):
             for n, param in previous_trainable_params.items():
                 new_param = trainer.model.get_parameter(n)
                 self.assertFalse(torch.allclose(param, new_param), f"Parameter {n} has not changed")
+
+    @parameterized.expand(
+        [
+            ("trl-internal-testing/tiny-Qwen2_5_VLForConditionalGeneration",),
+        ]
+    )
+    def test_train_vlm(self, model_id):
+        # Get the dataset
+        dataset = load_dataset("trl-internal-testing/zen-image", "conversational_language_modeling", split="train")
+
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            # Initialize the trainer
+            training_args = SFTConfig(output_dir=tmp_dir, report_to="none")
+            trainer = SFTTrainer(model=model_id, args=training_args, train_dataset=dataset)
+
+            # Save the initial parameters to compare them later
+            previous_trainable_params = {n: param.clone() for n, param in trainer.model.named_parameters()}
+
+            # Train the model
+            trainer.train()
+
+            # Check that the training loss is not None
+            self.assertIsNotNone(trainer.state.log_history[-1]["train_loss"])
+
+            # Check the params have changed
+            for n, param in previous_trainable_params.items():
+                new_param = trainer.model.get_parameter(n)
+                self.assertFalse(torch.allclose(param, new_param), f"Parameter {n} has not changed")

trl/trainer/grpo_trainer.py

@@ -462,7 +462,7 @@ 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`] or [`~transformers.ProcessorMixin`], *optional*, defaults to `None`):
+        processing_class ([`~transformers.PreTrainedTokenizerBase`] [`~transformers.ProcessorMixin`] or `None`, *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.AutoProcessor.from_pretrained`]. A
             padding token, `tokenizer.pad_token`, must be set. If the processing class has not set a padding token,
@@ -534,9 +534,9 @@ def __init__(
         else:
             model_id = model.config._name_or_path
             if args.model_init_kwargs is not None:
-                raise ValueError(
+                warnings.warn(
                     "You passed `model_init_kwargs` to the `GRPOConfig`, but your model is already instantiated. "
-                    "This argument can only be used when the `model` argument is a string."
+                    "The `model_init_kwargs` will be ignored."
                 )
 
         # Some models (SmolVLM/Idefics3) don't support `logits_to_keep` argument and error out if we pass it