[LLaVa] Add past_key_values to _skip_keys_device_placement to fix mul…

…ti-GPU dispatch (#28051) Add past_key_values to _skip_keys_device_placement for LLaVa
huggingface · Dec 15, 2023 · e2b6df7 · e2b6df7
1 parent deb72cb
commit e2b6df7
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 0 deletions.
diff --git a/src/transformers/models/llava/modeling_llava.py b/src/transformers/models/llava/modeling_llava.py
@@ -130,6 +130,7 @@ class LlavaPreTrainedModel(PreTrainedModel):
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["LlavaVisionAttention"]
+    _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = True
 
     def _init_weights(self, module):

diff --git a/src/transformers/models/vipllava/modeling_vipllava.py b/src/transformers/models/vipllava/modeling_vipllava.py
@@ -137,6 +137,7 @@ class VipLlavaPreTrainedModel(PreTrainedModel):
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
     _no_split_modules = ["VipLlavaVisionAttention"]
+    _skip_keys_device_placement = "past_key_values"
     _supports_flash_attn_2 = True
 
     def _init_weights(self, module):