[minor] phi4 train improvements (#1564)

Author: optas

configs/recipes/vision/phi4/sft/train.yaml

@@ -32,7 +32,7 @@ model:
 
 data:
   train:
-    collator_name: "vision_language_with_padding"
+    collator_name: "vision_language_sft"
     use_torchdata: true
     datasets:
       - dataset_name: "merve/vqav2-small"
@@ -71,7 +71,7 @@ training:
   output_dir: "output/vlm_finetuned"
   trainer_type: "TRL_SFT"
   enable_gradient_checkpointing: True
-  per_device_train_batch_size: 1 # Due to processor's handling of variable sized img-features.
+  per_device_train_batch_size: 2
   gradient_accumulation_steps: 8
   max_steps: 20
 

src/oumi/builders/collators.py

@@ -161,9 +161,10 @@ def build_collator_from_config(
         collator_kwargs["allow_multi_image_inputs"] = (
             model_config.visual_config.supports_multiple_images
         )
-        collator_kwargs["main_image_feature"] = (
-            model_config.visual_config.main_image_feature
-        )
+        if collator_name == "vision_language_with_padding":
+            collator_kwargs["main_image_feature"] = (
+                model_config.visual_config.main_image_feature
+            )
 
     if collator_name == "vision_language_sft":
         processor_name = collator_kwargs.get(

src/oumi/core/collators/vision_language_sft_collator.py

@@ -21,8 +21,6 @@
 from oumi.core.tokenizers.base_tokenizer import BaseTokenizer
 from oumi.core.types import Conversation
 
-_PIXEL_VALUES_KEY = "pixel_values"
-
 
 class VisionLanguageSftCollator:
     def __init__(

tests/unit/builders/test_collators.py

@@ -154,7 +154,7 @@ def test_build_collator_from_config_no_collator(mock_tokenizer):
     assert collator is None
 
 
-def test_build_collator_from_config_no_collator_no_tokenzier():
+def test_build_collator_from_config_no_collator_no_tokenizer():
     training_config = TrainingConfig(
         data=DataParams(
             train=DatasetSplitParams(