add vlm train factory

Author: RuBing-Yang

angelslim/compressor/speculative/__init__.py

@@ -14,10 +14,9 @@
 
 from .benchmark import BenchmarkConfig, BenchmarkEngine, BenchmarkMode
 from .train import (
-    DataCollatorWithPadding,
     DatasetManager,
     DraftModelConfig,
-    OnlineEagle3Trainer,
+    Eagle3TrainerFactory,
     convert_sharegpt_data,
     convert_ultrachat_data,
     create_draft_model,
@@ -33,9 +32,8 @@
     "create_draft_model",
     "DraftModelConfig",
     "create_target_model",
-    "OnlineEagle3Trainer",
+    "Eagle3TrainerFactory",
     "data_generation_work_flow",
-    "DataCollatorWithPadding",
     "convert_sharegpt_data",
     "convert_ultrachat_data",
     "DatasetManager",

angelslim/compressor/speculative/train/__init__.py

@@ -1,21 +1,19 @@
 from .data import (
-    DataCollatorWithPadding,
     DatasetManager,
     convert_sharegpt_data,
     convert_ultrachat_data,
     data_generation_work_flow,
     get_supported_chat_template_type_strings,
 )
 from .models import DraftModelConfig, create_draft_model, create_target_model
-from .trainer import OnlineEagle3Trainer
+from .trainer import Eagle3TrainerFactory
 
 __all__ = [
     "create_draft_model",
     "DraftModelConfig",
     "create_target_model",
-    "OnlineEagle3Trainer",
+    "Eagle3TrainerFactory",
     "data_generation_work_flow",
-    "DataCollatorWithPadding",
     "convert_sharegpt_data",
     "convert_ultrachat_data",
     "DatasetManager",

angelslim/compressor/speculative/train/data/__init__.py

@@ -14,16 +14,11 @@
 
 from .chat_templates import get_supported_chat_template_type_strings
 from .data_generation import data_generation_work_flow
-from .data_utils import (
-    DataCollatorWithPadding,
-    convert_sharegpt_data,
-    convert_ultrachat_data,
-)
+from .data_utils import convert_sharegpt_data, convert_ultrachat_data
 from .dataset import DatasetManager
 
 __all__ = [
     "DatasetManager",
-    "DataCollatorWithPadding",
     "convert_sharegpt_data",
     "convert_ultrachat_data",
     "data_generation_work_flow",

angelslim/compressor/speculative/train/data/data_utils.py

@@ -21,6 +21,7 @@
     "convert_sharegpt_data",
     "convert_ultrachat_data",
     "DataCollatorWithPadding",
+    "VLMDataCollatorWithPadding",
 ]
 
 
@@ -100,38 +101,92 @@ def process_token_dict_to_mappings(
     return d2t, t2d
 
 
+def paddingtensor(intensors, N):
+    B, n, S = intensors.shape
+    # padding_tensor = torch.zeros(B, N - n, S,dtype=intensors.dtype)
+    padding_tensor = torch.zeros(B, N - n, S, dtype=intensors.dtype)
+    outtensors = torch.cat((intensors, padding_tensor), dim=1)
+    return outtensors
+
+
+def paddingtensor2D(intensors, N):
+    B, n = intensors.shape
+    padding_tensor = torch.zeros(B, N - n, dtype=intensors.dtype)
+    outtensors = torch.cat((intensors, padding_tensor), dim=1)
+    return outtensors
+
+
+def paddingtensor3D(tensor_list):
+    max_h = max(tensor.shape[-2] for tensor in tensor_list)
+    max_w = max(tensor.shape[-1] for tensor in tensor_list)
+    out_tensor_list = []
+    for tensor in tensor_list:
+        if tensor.ndim == 2:
+            tensor = tensor.unsqueeze(0)
+        b, h, w = tensor.shape
+        outtensor = torch.zeros(b, max_h, max_w, dtype=tensor.dtype)
+        outtensor[:, :h, :w] = tensor
+        out_tensor_list.append(outtensor)
+    return torch.cat(out_tensor_list)
+
+
 class DataCollatorWithPadding:
-    def paddingtensor(self, intensors, N):
-        B, n, S = intensors.shape
-        # padding_tensor = torch.zeros(B, N - n, S,dtype=intensors.dtype)
-        padding_tensor = torch.zeros(B, N - n, S, dtype=intensors.dtype)
-        outtensors = torch.cat((intensors, padding_tensor), dim=1)
-        return outtensors
-
-    def paddingtensor2D(self, intensors, N):
-        B, n = intensors.shape
-        padding_tensor = torch.zeros(B, N - n, dtype=intensors.dtype)
-        outtensors = torch.cat((intensors, padding_tensor), dim=1)
-        return outtensors
 
     def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
         max_length = max(item["input_ids"].shape[1] for item in features)
         batch_input_ids = torch.cat(
-            [self.paddingtensor2D(item["input_ids"], max_length) for item in features]
+            [paddingtensor2D(item["input_ids"], max_length) for item in features]
+        )
+        batch_attention_mask = torch.cat(
+            [paddingtensor2D(item["attention_mask"], max_length) for item in features]
+        )
+        batch_loss_mask = torch.cat(
+            [paddingtensor2D(item["loss_mask"], max_length) for item in features]
+        )
+
+        batch = {
+            "input_ids": batch_input_ids,
+            "attention_mask": batch_attention_mask,
+            "loss_mask": batch_loss_mask,
+        }
+        return batch
+
+
+class VLMDataCollatorWithPadding:
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        max_length = max(item["input_ids"].shape[1] for item in features)
+        batch_input_ids = torch.cat(
+            [paddingtensor2D(item["input_ids"], max_length) for item in features]
         )
         batch_attention_mask = torch.cat(
-            [
-                self.paddingtensor2D(item["attention_mask"], max_length)
-                for item in features
-            ]
+            [paddingtensor2D(item["attention_mask"], max_length) for item in features]
         )
         batch_loss_mask = torch.cat(
-            [self.paddingtensor2D(item["loss_mask"], max_length) for item in features]
+            [paddingtensor2D(item["loss_mask"], max_length) for item in features]
         )
 
         batch = {
             "input_ids": batch_input_ids,
             "attention_mask": batch_attention_mask,
             "loss_mask": batch_loss_mask,
         }
+
+        if "pixel_values" in features[0]:
+            batch["pixel_values"] = paddingtensor3D(
+                [item["pixel_values"] for item in features]
+            )
+        if "video_pixel_values" in features[0]:
+            batch["video_pixel_values"] = paddingtensor3D(
+                [item["video_pixel_values"] for item in features]
+            )
+        if "image_grid_thw" in features[0]:
+            batch["image_grid_thw"] = paddingtensor3D(
+                [item["image_grid_thw"] for item in features]
+            )
+        if "video_grid_thw" in features[0]:
+            batch["video_grid_thw"] = paddingtensor3D(
+                [item["video_grid_thw"] for item in features]
+            )
+
         return batch