Update modeling_internvl_chat.py

modeling_internvl_chat.py

@@ -20,8 +20,8 @@ from .configuration_internvl_chat import InternVLChatConfig
 from .modeling_intern_vit import InternVisionModel
 from .modeling_phi3 import Phi3ForCausalLM
 
-logger = logging.get_logger(__name__)
 
+logger = logging.get_logger(__name__)
 
 class InternVLChatModel(PreTrainedModel):
     config_class = InternVLChatConfig
@@ -31,6 +31,7 @@ class InternVLChatModel(PreTrainedModel):
     def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
         super().__init__(config)
 
+        # Initialize components
         image_size = config.force_image_size or config.vision_config.image_size
         patch_size = config.vision_config.patch_size
         self.patch_size = patch_size
@@ -42,10 +43,12 @@ class InternVLChatModel(PreTrainedModel):
 
         logger.info(f'num_image_token: {self.num_image_token}')
         logger.info(f'ps_version: {self.ps_version}')
+
         if vision_model is not None:
             self.vision_model = vision_model
         else:
             self.vision_model = InternVisionModel(config.vision_config)
+
         if language_model is not None:
             self.language_model = language_model
         else:
@@ -56,6 +59,11 @@ class InternVLChatModel(PreTrainedModel):
             else:
                 raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
 
+        # Wrap models with DataParallel for multi-GPU support
+        if torch.cuda.device_count() > 1:
+            self.vision_model = nn.DataParallel(self.vision_model)
+            self.language_model = nn.DataParallel(self.language_model)
+
         vit_hidden_size = config.vision_config.hidden_size
         llm_hidden_size = config.llm_config.hidden_size
 
@@ -66,13 +74,6 @@ class InternVLChatModel(PreTrainedModel):
             nn.Linear(llm_hidden_size, llm_hidden_size)
         )
 
-        # if config.force_image_size != config.vision_config.image_size:
-        #     self.vision_model.resize_pos_embeddings(
-        #         old_size=config.vision_config.image_size,
-        #         new_size=config.force_image_size,
-        #         patch_size=config.vision_config.patch_size
-        #     )
-
         self.img_context_token_id = None
         self.neftune_alpha = None
 
@@ -82,6 +83,7 @@ class InternVLChatModel(PreTrainedModel):
         if config.use_llm_lora:
             self.wrap_llm_lora(r=config.use_llm_lora, lora_alpha=2 * config.use_llm_lora)
 
+
     def wrap_backbone_lora(self, r=128, lora_alpha=256, lora_dropout=0.05):
         lora_config = LoraConfig(
             r=r,