Add vision_batch_size to avoid cuda OOM (#4)

- Add vision_batch_size to avoid cuda OOM (74e66369a97903f4922eaaa1f0c6d5d5a591faf9)

configuration_minicpm.py

@@ -69,6 +69,7 @@ class MiniCPMVConfig(Qwen2Config):
         slice_config=None,
         vision_config=None,
         use_image_id=True,
+        vision_batch_size=16,
         **kwargs,
     ):
         self.use_cache = use_cache
@@ -77,6 +78,7 @@ class MiniCPMVConfig(Qwen2Config):
         self.drop_vision_last_layer = drop_vision_last_layer
         self.batch_vision_input = batch_vision_input
         self.use_image_id = use_image_id
+        self.vision_batch_size = vision_batch_size
 
         if slice_config is None:
             self.slice_config = MiniCPMVSliceConfig(max_slice_nums=1)
@@ -95,4 +97,4 @@ class MiniCPMVConfig(Qwen2Config):
 
         self.patch_size = self.vision_config.patch_size
 
-        super().__init__(**kwargs)
+        super().__init__(**kwargs)

modeling_minicpmv.py

@@ -92,31 +92,30 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 tgt_sizes = [tgt_size for tgt_size in tgt_sizes if isinstance(tgt_size, torch.Tensor)]
                 tgt_sizes = torch.vstack(tgt_sizes).type(torch.int32)
 
-                if self.config.batch_vision_input:
-                    max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
-
-                    all_pixel_values = torch.nn.utils.rnn.pad_sequence(all_pixel_values, batch_first=True,
-                                                                       padding_value=0.0)
-                    B, L, _ = all_pixel_values.shape
-                    all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
-
-                    patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device)
-                    for i in range(B):
-                        patch_attn_mask[i, 0, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True
-
-                    vision_embedding = self.vpm(all_pixel_values.type(dtype), patch_attention_mask=patch_attn_mask, tgt_sizes=tgt_sizes).last_hidden_state
-                    vision_embedding = self.resampler(vision_embedding, tgt_sizes)
+                max_patches = torch.max(tgt_sizes[:, 0] * tgt_sizes[:, 1])
+
+                all_pixel_values = torch.nn.utils.rnn.pad_sequence(all_pixel_values, batch_first=True,
+                                                                   padding_value=0.0)
+                B, L, _ = all_pixel_values.shape
+                all_pixel_values = all_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
+
+                patch_attn_mask = torch.zeros((B, 1, max_patches), dtype=torch.bool, device=device)
+                for i in range(B):
+                    patch_attn_mask[i, 0, :tgt_sizes[i][0] * tgt_sizes[i][1]] = True
+
+                vision_batch_size = self.config.vision_batch_size
+                all_pixel_values = all_pixel_values.type(dtype)
+                if B > vision_batch_size:
+                    hs = []
+                    for i in range(0, B, vision_batch_size):
+                        start_idx = i
+                        end_idx = i + vision_batch_size
+                        tmp_hs = self.vpm(all_pixel_values[start_idx:end_idx], patch_attention_mask=patch_attn_mask[start_idx:end_idx], tgt_sizes=tgt_sizes[start_idx:end_idx]).last_hidden_state
+                        hs.append(tmp_hs)
+                    vision_embedding = torch.cat(hs, dim=0)
                 else:
-                    # get vision_embedding foreach
-                    vision_embedding = []
-                    for single_tgt_size, single_pixel_values in zip(tgt_sizes, all_pixel_values):
-                        single_pixel_values = single_pixel_values.unsqueeze(0)
-                        B, L, _ = single_pixel_values.shape
-                        single_pixel_values = single_pixel_values.permute(0, 2, 1).reshape(B, 3, -1, L)
-                        single_vision_embedding = self.vpm(single_pixel_values.type(dtype), tgt_sizes=single_tgt_size.unsqueeze(0)).last_hidden_state
-                        single_vision_embedding = self.resampler(single_vision_embedding, single_tgt_size.unsqueeze(0))
-                        vision_embedding.append(single_vision_embedding)
-                    vision_embedding = torch.vstack(vision_embedding)
+                    vision_embedding = self.vpm(all_pixel_values, patch_attention_mask=patch_attn_mask, tgt_sizes=tgt_sizes).last_hidden_state
+                vision_embedding = self.resampler(vision_embedding, tgt_sizes)
 
                 start = 0
                 for pixel_values in pixel_values_list: