Add supports_gradient_checkpointing

This commit is contained in:
ai-modelscope
2025-03-26 13:52:52 +08:00
parent d13495871e
commit b2bedcd779
3 changed files with 14 additions and 0 deletions

View File

@@ -64,6 +64,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}')

View File

@@ -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']

View File

@@ -38,6 +38,7 @@ 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', 'Phi3DecoderLayer']
def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None, use_flash_attn=True):
@@ -347,3 +348,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()