Support device_map='auto'

OpenGVLab · Feb 4, 2024 · 21baf4c · 21baf4c
1 parent 5aa95de
commit 21baf4c
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 1 deletion.
diff --git a/internvl_chat/internvl/model/internvl_chat/modeling_intern_vit.py b/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)

diff --git a/internvl_chat/internvl/model/internvl_chat/modeling_internvl_chat.py b/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)

diff --git a/internvl_chat/internvl/model/internvl_chat_with_qllama/modeling_intern_vit.py b/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)

diff --git a/internvl_chat/internvl/model/internvl_chat_with_qllama/modeling_internvl.py b/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']
 

diff --git a/internvl_chat/internvl/model/internvl_chat_with_qllama/modeling_internvl_chat.py b/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)