Upload folder using huggingface_hub
Browse files
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}')
|