[VLM] Multimodal Data Collator (#1087)

## Purpose ##
* Move data collators to example script as per @mgoin's suggestion
#1032 (comment)

## Changes ##
* Remove data collator definitions in LC
* Add data collators in examples with a comment indicating that this is
for multimodal inputs

## Testing ##
Ran all multimodal vision models
* Qwen2
* Pixtral
* Mllama
* Llava
* Phi3_vision

---------

Signed-off-by: Kyle Sayers <[email protected]>

Author: kylesayrs

examples/multimodal_vision/llava_example.py

@@ -1,11 +1,11 @@
 import requests
+import torch
 from PIL import Image
 from transformers import AutoProcessor
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.transformers import oneshot
 from llmcompressor.transformers.tracing import TraceableLlavaForConditionalGeneration
-from llmcompressor.transformers.utils.data_collator import llava_data_collator
 
 # Load model.
 model_id = "llava-hf/llava-1.5-7b-hf"
@@ -20,6 +20,13 @@
 NUM_CALIBRATION_SAMPLES = 512
 MAX_SEQUENCE_LENGTH = 2048
 
+
+# Define a oneshot data collator for multimodal inputs.
+def data_collator(batch):
+    assert len(batch) == 1
+    return {key: torch.tensor(value) for key, value in batch[0].items()}
+
+
 # Recipe
 recipe = [
     GPTQModifier(
@@ -40,7 +47,7 @@
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=llava_data_collator,
+    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.

examples/multimodal_vision/mllama_example.py

@@ -1,11 +1,11 @@
 import requests
+import torch
 from PIL import Image
 from transformers import AutoProcessor
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.transformers import oneshot
 from llmcompressor.transformers.tracing import TraceableMllamaForConditionalGeneration
-from llmcompressor.transformers.utils.data_collator import mllama_data_collator
 
 # Load model.
 model_id = "meta-llama/Llama-3.2-11B-Vision-Instruct"
@@ -20,6 +20,13 @@
 NUM_CALIBRATION_SAMPLES = 512
 MAX_SEQUENCE_LENGTH = 2048
 
+
+# Define a oneshot data collator for multimodal inputs.
+def data_collator(batch):
+    assert len(batch) == 1
+    return {key: torch.tensor(value) for key, value in batch[0].items()}
+
+
 # Recipe
 recipe = [
     GPTQModifier(
@@ -39,7 +46,7 @@
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=mllama_data_collator,
+    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.

examples/multimodal_vision/phi3_vision_example.py

@@ -1,10 +1,10 @@
+import torch
 from datasets import load_dataset
 from transformers import AutoModelForCausalLM, AutoProcessor
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
 from llmcompressor.transformers import oneshot
-from llmcompressor.transformers.utils.data_collator import phi3_vision_data_collator
 
 # Load model.
 model_id = "microsoft/Phi-3-vision-128k-instruct"
@@ -60,6 +60,12 @@ def tokenize(sample):
 ds = ds.map(tokenize, writer_batch_size=1, remove_columns=ds.column_names)
 
 
+# Define a oneshot data collator for multimodal inputs.
+def data_collator(batch):
+    assert len(batch) == 1
+    return {key: torch.tensor(value) for key, value in batch[0].items()}
+
+
 # Recipe
 recipe = [
     SmoothQuantModifier(smoothing_strength=0.8),
@@ -79,7 +85,7 @@ def tokenize(sample):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=phi3_vision_data_collator,
+    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.

examples/multimodal_vision/pixtral_example.py

@@ -1,11 +1,11 @@
 import requests
+import torch
 from PIL import Image
 from transformers import AutoProcessor
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.transformers import oneshot
 from llmcompressor.transformers.tracing import TraceableLlavaForConditionalGeneration
-from llmcompressor.transformers.utils.data_collator import pixtral_data_collator
 
 # Load model.
 model_id = "mgoin/pixtral-12b"
@@ -20,6 +20,17 @@
 NUM_CALIBRATION_SAMPLES = 512
 MAX_SEQUENCE_LENGTH = 2048
 
+
+# Define a oneshot data collator for multimodal inputs.
+def data_collator(batch):
+    assert len(batch) == 1
+    return {
+        "input_ids": torch.LongTensor(batch[0]["input_ids"]),
+        "attention_mask": torch.tensor(batch[0]["attention_mask"]),
+        "pixel_values": torch.tensor(batch[0]["pixel_values"])[0],
+    }
+
+
 # Recipe
 recipe = [
     GPTQModifier(
@@ -40,7 +51,7 @@
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=pixtral_data_collator,
+    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.

examples/multimodal_vision/qwen2_vl_example.py

@@ -1,14 +1,14 @@
 import base64
 from io import BytesIO
 
+import torch
 from datasets import load_dataset
 from qwen_vl_utils import process_vision_info
 from transformers import AutoProcessor
 
 from llmcompressor.modifiers.quantization import GPTQModifier
 from llmcompressor.transformers import oneshot
 from llmcompressor.transformers.tracing import TraceableQwen2VLForConditionalGeneration
-from llmcompressor.transformers.utils.data_collator import qwen2_vl_data_collator
 
 # Load model.
 model_id = "Qwen/Qwen2-VL-2B-Instruct"
@@ -65,6 +65,13 @@ def preprocess_and_tokenize(example):
 
 ds = ds.map(preprocess_and_tokenize, remove_columns=ds["calibration"].column_names)
 
+
+# Define a oneshot data collator for multimodal inputs.
+def data_collator(batch):
+    assert len(batch) == 1
+    return {key: torch.tensor(value) for key, value in batch[0].items()}
+
+
 # Recipe
 recipe = [
     GPTQModifier(
@@ -84,7 +91,7 @@ def preprocess_and_tokenize(example):
     max_seq_length=MAX_SEQUENCE_LENGTH,
     num_calibration_samples=NUM_CALIBRATION_SAMPLES,
     trust_remote_code_model=True,
-    data_collator=qwen2_vl_data_collator,
+    data_collator=data_collator,
 )
 
 # Confirm generations of the quantized model look sane.