fix MultiModalDataset padding

Author: irisliu10

angelslim/data/dataloader.py

@@ -43,6 +43,7 @@ def create_data_loader(
         inference_settings: Dict = None,
         use_audio_in_video: bool = False,
         model_name: str = None,
+        quantization_config: str = None,
     ) -> DataLoader:
         """
         Create appropriate DataLoader based on data source
@@ -94,6 +95,7 @@ def create_data_loader(
                 data_source=data_source,
                 is_hf_dataset=not os.path.isfile(data_source),
                 model_name=model_name,
+                quantization_config=quantization_config,
             )
         elif data_type == "Text2ImageDataset":
             dataset = Text2ImageDataset(

angelslim/data/multimodal_dataset.py

@@ -37,10 +37,12 @@ def __init__(
         data_source: Union[str, Dict] = None,
         is_hf_dataset: bool = False,
         model_name: str = None,
+        quantization_config: str = None,
     ):
         super().__init__(processor, device, max_length)
         self.is_hf_dataset = is_hf_dataset
         self.model_name = model_name
+        self.quant_algo = quantization_config.name if quantization_config else None
 
         if is_hf_dataset:
             self._load_hf_dataset(data_source, num_samples)
@@ -174,13 +176,21 @@ def _load_hf_dataset(self, dataset: str, num_samples: int):
 
     def _process_and_append(self, messages: List[Dict], tools=None):
         """Process messages and append to dataset"""
+
+        # max_length padding for gptq and awq
+        if "gptq" in self.quant_algo or "awq" in self.quant_algo:
+            padding = "max_length"
+        else:
+            padding = True
+
         if self.model_name in ["Qwen3VL", "Qwen3VLMoE"]:
             inputs = self.processor.apply_chat_template(
                 messages,
                 tools=tools,
                 tokenize=True,
                 add_generation_prompt=True,
                 return_dict=True,
+                padding=padding,
                 truncation=True,
                 return_tensors="pt",
                 max_length=self.max_length,
@@ -195,6 +205,7 @@ def _process_and_append(self, messages: List[Dict], tools=None):
             inputs = self.processor(
                 text=[text],
                 images=image_inputs,
+                padding=padding,
                 truncation=True,
                 return_tensors="pt",
                 max_length=self.max_length,
@@ -212,6 +223,7 @@ def _process_and_append(self, messages: List[Dict], tools=None):
                 text=[text],
                 images=image_inputs,
                 videos=video_inputs,
+                padding=padding,
                 truncation=True,
                 return_tensors="pt",
                 max_length=self.max_length,

angelslim/engine.py

@@ -149,6 +149,7 @@ def prepare_data(
         inference_settings=None,
         use_audio_in_video=False,
         model_name=None,
+        quantization_config=None,
     ) -> Optional[Any]:
         """Prepare compression dataset"""
         if custom_dataloader is not None:
@@ -174,6 +175,7 @@ def prepare_data(
             inference_settings=inference_settings,
             use_audio_in_video=use_audio_in_video,
             model_name=model_name,
+            quantization_config=quantization_config,
         )
         self.max_seq_length = max_length
 

tools/run.py

@@ -169,6 +169,7 @@ def run(config):
             inference_settings=dataset_config.inference_settings,
             use_audio_in_video=model_config.use_audio_in_video,
             model_name=model_config.name,
+            quantization_config=compress_config.quantization,
         )
 
     # Step 5: Initialize compressor