Upload folder using huggingface_hub

modeling_internvl_chat.py

@@ -38,7 +38,7 @@ class InternVLChatModel(PreTrainedModel):
     _supports_flash_attn_2 = True
     _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer', 'InternLM2DecoderLayer']
 
-    def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
+    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')
@@ -50,6 +50,8 @@ class InternVLChatModel(PreTrainedModel):
         self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
         self.downsample_ratio = config.downsample_ratio
         self.ps_version = config.ps_version
+        config.vision_config.use_flash_attn = True if use_flash_attn else False
+        config.llm_config.attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
 
         logger.info(f'num_image_token: {self.num_image_token}')
         logger.info(f'ps_version: {self.ps_version}')