Update model code

InternVL2-40B-Pretrain/configuration_internvl_chat.py

@@ -61,6 +61,8 @@ class InternVLChatConfig(PretrainedConfig):
         self.ps_version = ps_version  # pixel shuffle version
         self.min_dynamic_patch = min_dynamic_patch
         self.max_dynamic_patch = max_dynamic_patch
+        # By default, we use tie_word_embeddings=False for models of all sizes.
+        self.tie_word_embeddings = self.llm_config.tie_word_embeddings
 
         logger.info(f'vision_select_layer: {self.select_layer}')
         logger.info(f'ps_version: {self.ps_version}')

InternVL2-40B-Pretrain/modeling_intern_vit.py

@@ -364,6 +364,7 @@ class InternVisionEncoder(nn.Module):
 class InternVisionModel(PreTrainedModel):
     main_input_name = 'pixel_values'
     _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
     config_class = InternVisionConfig
     _no_split_modules = ['InternVisionEncoderLayer']
 

InternVL2-40B-Pretrain/modeling_internvl_chat.py

@@ -36,12 +36,13 @@ class InternVLChatModel(PreTrainedModel):
     main_input_name = 'pixel_values'
     base_model_prefix = 'language_model'
     _supports_flash_attn_2 = True
+    supports_gradient_checkpointing = True
     _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer']
 
     def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
         super().__init__(config)
 
-        assert version_cmp(transformers.__version__, '4.36.2', 'ge')
+        assert version_cmp(transformers.__version__, '4.37.0', 'ge')
         image_size = config.force_image_size or config.vision_config.image_size
         patch_size = config.vision_config.patch_size
         self.patch_size = patch_size
@@ -108,7 +109,7 @@ class InternVLChatModel(PreTrainedModel):
         B, N, C = input_embeds.shape
         input_embeds = input_embeds.reshape(B * N, C)
 
-        if torch.distributed.get_rank() == 0:
+        if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:
             print(f'dynamic ViT batch size: {vit_batch_size}, images per sample: {vit_batch_size / B}, dynamic token length: {N}')
 
         input_ids = input_ids.reshape(B * N)
@@ -343,3 +344,13 @@ class InternVLChatModel(PreTrainedModel):
         )
 
         return outputs
+
+    @property
+    def lm_head(self):
+        return self.language_model.get_output_embeddings()
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.language_model.get_output_embeddings()