Update modeling_minicpmv.py

update model.py for zero3 and lora finetuning

modeling_minicpmv.py

@@ -1,17 +1,20 @@
 import math
-from typing import List, Optional
 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
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from peft.utils.other import ModulesToSaveWrapper
 
 
 class MiniCPMVPreTrainedModel(MiniCPMPreTrainedModel):
@@ -72,17 +75,35 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
     def set_input_embeddings(self, value):
         self.llm.embed_tokens = value
 
+    def vpm_forward_features(self, pixel_value):
+        if isinstance(self.vpm, ModulesToSaveWrapper):
+            if self.vpm.disable_adapters or (self.vpm.active_adapter not in self.vpm.modules_to_save):
+                return self.vpm.original_module.forward_features(pixel_value)
+            return self.vpm.modules_to_save[self.vpm.active_adapter].forward_features(pixel_value)
+        else:    
+            return self.vpm.forward_features(pixel_value)
+            
     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]))
-            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))
+            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(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):
+                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:
+                print(pixel_value.shape)
+                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 +114,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
                     )
@@ -319,24 +340,21 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
             content = msg["content"]
             assert role in ["user", "assistant"]
             if i == 0:
-                if image is None:
-                    images = []
+                assert role == "user", "The role of first msg should be user"
+                if self.config.slice_mode:
+                    images, final_placeholder = self.get_slice_image_placeholder(
+                        image, tokenizer
+                    )
+                    content = final_placeholder + "\n" + content
                 else:
-                    assert role == "user", "The role of first msg should be user"
-                    if self.config.slice_mode:
-                        images, final_placeholder = self.get_slice_image_placeholder(
-                            image, tokenizer
-                        )
-                        content = final_placeholder + "\n" + content
-                    else:
-                        images = [image]
-                        content = (
-                            tokenizer.im_start
-                            + tokenizer.unk_token * self.config.query_num
-                            + tokenizer.im_end
-                            + "\n"
-                            + content
-                        )
+                    images = [image]
+                    content = (
+                        tokenizer.im_start
+                        + tokenizer.unk_token * self.config.query_num
+                        + tokenizer.im_end
+                        + "\n"
+                        + content
+                    )
             prompt += "<用户>" if role == "user" else "<AI>"
             prompt += content
         prompt += "<AI>"
@@ -377,6 +395,8 @@ class MiniCPMV(MiniCPMVPreTrainedModel):
 
         return answer, context, generation_config
 
+    
+
 
 class LlamaTokenizerWrapper(LlamaTokenizer):
     def __init__(self, **kwargs):