Update modeling_minicpmv.py

modeling_minicpmv.py

@@ -4,11 +4,12 @@ import json
 import timm
 import torch
 import torchvision
+import deepspeed
 from PIL import Image
 from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
 from torchvision import transforms
 from transformers import LlamaTokenizer
-
+from transformers.integrations import is_deepspeed_zero3_enabled
 from .configuration_minicpm import MiniCPMVConfig
 from .modeling_minicpm import MiniCPMForCausalLM, MiniCPMPreTrainedModel
 from .resampler import Resampler
@@ -74,15 +75,24 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
     def get_vision_embedding(self, pixel_values):
         res = []
-        dtype = self.vpm.pos_embed.data.dtype
-        for pixel_value in pixel_values:
+        dtype = self.llm.lm_head.weight.dtype
+        def process_each_pixel(pixel_value, dtype, config, vpm, resampler):
             H, W = pixel_value.shape[-2:]
-            tgt_size = (
-            math.ceil(H / self.vpm.patch_embed.patch_size[0]), math.ceil(W / self.vpm.patch_embed.patch_size[0]))
+            target_size = (math.ceil(H / config.patch_size), math.ceil(W / config.patch_size))
             vision_embedding = self.vpm.forward_features(pixel_value.unsqueeze(0).type(dtype))
-            if hasattr(self.vpm, 'num_prefix_tokens') and self.vpm.num_prefix_tokens > 0:
-                vision_embedding = vision_embedding[:, self.vpm.num_prefix_tokens:]
-            res.append(self.resampler(vision_embedding, tgt_size))
+            if hasattr(vpm, 'num_prefix_tokens') and vpm.num_prefix_tokens > 0:
+                vision_embedding = vision_embedding[:, vpm.num_prefix_tokens:]
+            return resampler(vision_embedding, target_size)
+
+        if is_deepspeed_zero3_enabled():
+            with deepspeed.zero.GatheredParameters(self.vpm.pos_embed, modifier_rank=0):
+                for pixel_value in pixel_values:
+                    result = process_each_pixel(pixel_value, dtype, self.config, self.vpm, self.resampler)
+                    res.append(result)
+        else:
+            for pixel_value in pixel_values:
+                result = process_each_pixel(pixel_value, dtype, self.config, self.vpm, self.resampler)
+                res.append(result)
         return torch.vstack(res)
 
     def get_vllm_embedding(self, data):
@@ -93,8 +103,8 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
                 if len(pixel_values) > 0:
                     vision_hidden_states.append(self.get_vision_embedding(pixel_values))
                 elif self.training:
-                    dtype = self.vpm.pos_embed.data.dtype
-                    device = self.vpm.pos_embed.data.device
+                    dtype = self.llm.lm_head.weight.dtype
+                    device = self.llm.lm_head.weight.device
                     dummy_image = torch.zeros(
                         (1, 3, 224, 224), device=device, dtype=dtype
                     )