czczup commited on
Commit
b038051
1 Parent(s): 980f3ea

Upload folder using huggingface_hub

Browse files
Files changed (1) hide show
  1. modeling_internvl_chat.py +3 -1
modeling_internvl_chat.py CHANGED
@@ -36,7 +36,7 @@ class InternVLChatModel(PreTrainedModel):
36
  _supports_flash_attn_2 = True
37
  _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer']
38
 
39
- def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
40
  super().__init__(config)
41
 
42
  assert version_cmp(transformers.__version__, '4.36.2', 'ge')
@@ -48,6 +48,8 @@ class InternVLChatModel(PreTrainedModel):
48
  self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
49
  self.downsample_ratio = config.downsample_ratio
50
  self.ps_version = config.ps_version
 
 
51
 
52
  logger.info(f'num_image_token: {self.num_image_token}')
53
  logger.info(f'ps_version: {self.ps_version}')
 
36
  _supports_flash_attn_2 = True
37
  _no_split_modules = ['InternVisionModel', 'LlamaDecoderLayer']
38
 
39
+ def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
40
  super().__init__(config)
41
 
42
  assert version_cmp(transformers.__version__, '4.36.2', 'ge')
 
48
  self.num_image_token = int((image_size // patch_size) ** 2 * (config.downsample_ratio ** 2))
49
  self.downsample_ratio = config.downsample_ratio
50
  self.ps_version = config.ps_version
51
+ config.vision_config.use_flash_attn = True if use_flash_attn else False
52
+ config.llm_config._attn_implementation = 'flash_attention_2' if use_flash_attn else 'eager'
53
 
54
  logger.info(f'num_image_token: {self.num_image_token}')
55
  logger.info(f'ps_version: {self.ps_version}')