Skip to content

GitLab

Unverified Commit 3802e786 authored by Darshana S's avatar Darshana S Committed by GitHub
Browse files

Enable device map (#30870) 

* added_no_split_modules

* added LlavaNextVisionAttention to _no_split_modules
parent 57c965a8
Hide whitespace changes
Inline Side-by-side
Showing with 1 addition and 0 deletions
src/transformers/models/video_llava/modeling_video_llava.py
View file @ 3802e786
...@@ -124,6 +124,7 @@ class VideoLlavaPreTrainedModel(PreTrainedModel): ...@@ -124,6 +124,7 @@ class VideoLlavaPreTrainedModel(PreTrainedModel):
supports_gradient_checkpointing = True supports_gradient_checkpointing = True
_skip_keys_device_placement = "past_key_values" _skip_keys_device_placement = "past_key_values"
_supports_flash_attn_2 = True _supports_flash_attn_2 = True
_no_split_modules = ["VideoLlavaVisionAttention"]
def _init_weights(self, module): def _init_weights(self, module):
# important: this ported version of VideoLlava isn't meant for training from scratch - only # important: this ported version of VideoLlava isn't meant for training from scratch - only
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment