Support device_map='auto'

Author: czczup

internvl_chat/internvl/model/internvl_chat/modeling_intern_vit.py

@@ -279,6 +279,7 @@ def forward(
 class InternVisionModel(PreTrainedModel):
     main_input_name = 'pixel_values'
     config_class = InternVisionConfig
+    _no_split_modules = ['InternVisionEncoderLayer']
 
     def __init__(self, config: InternVisionConfig):
         super().__init__(config)

internvl_chat/internvl/model/internvl_chat/modeling_internvl_chat.py

@@ -23,6 +23,7 @@
 class InternVLChatModel(PreTrainedModel):
     config_class = InternVLChatConfig
     main_input_name = 'pixel_values'
+    _no_split_modules = ['InternVisionEncoderLayer', 'LlamaDecoderLayer', 'LlamaForCausalLM']
 
     def __init__(self, config: InternVLChatConfig, vision_model=None, language_model=None):
         super().__init__(config)

internvl_chat/internvl/model/internvl_chat_with_qllama/modeling_intern_vit.py

@@ -279,6 +279,7 @@ def forward(
 class InternVisionModel(PreTrainedModel):
     main_input_name = 'pixel_values'
     config_class = InternVisionConfig
+    _no_split_modules = ['InternVisionEncoderLayer']
 
     def __init__(self, config: InternVisionConfig):
         super().__init__(config)

internvl_chat/internvl/model/internvl_chat_with_qllama/modeling_internvl.py

@@ -42,7 +42,7 @@ class InternVLPreTrainedModel(PreTrainedModel):
     _keys_to_ignore_on_load_missing = [
         r'position_ids',
     ]
-    _no_split_modules = ['InternAttention', 'LlamaDecoderLayer', 'LlamaForCausalLM']
+    _no_split_modules = ['InternVisionEncoderLayer', 'LlamaDecoderLayer', 'LlamaForCausalLM']
     _skip_keys_device_placement = 'past_key_values'
     _keep_in_fp32_modules = ['wo']
 

internvl_chat/internvl/model/internvl_chat_with_qllama/modeling_internvl_chat.py

@@ -24,6 +24,7 @@
 class InternVLChatModel(PreTrainedModel):
     config_class = InternVLChatConfig
     main_input_name = 'pixel_values'
+    _no_split_modules = ['InternVisionEncoderLayer', 'LlamaDecoderLayer', 'LlamaForCausalLM']
 
     def __init__(self, config: InternVLChatConfig, internvl=None, language_model=None):
         super().__init__(config)