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
......@@ -124,6 +124,7 @@ class VideoLlavaPreTrainedModel(PreTrainedModel):
supports_gradient_checkpointing = True
_skip_keys_device_placement = "past_key_values"
_supports_flash_attn_2 = True
_no_split_modules = ["VideoLlavaVisionAttention"]
def _init_weights(self, module):
# 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