diff --git a/src/transformers/models/llava_next/modeling_llava_next.py b/src/transformers/models/llava_next/modeling_llava_next.py index 718c1e96c988..71d87501dbd5 100644 --- a/src/transformers/models/llava_next/modeling_llava_next.py +++ b/src/transformers/models/llava_next/modeling_llava_next.py @@ -32,7 +32,8 @@ from ...utils import LossKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging from ..auto import AutoModel from .configuration_llava_next import LlavaNextConfig - +from ..llama.modeling_llama import LlamaDecoderLayer +from ..clip.modeling_clip import CLIPEncoderLayer logger = logging.get_logger(__name__) @@ -245,7 +246,7 @@ class LlavaNextPreTrainedModel(PreTrainedModel): config_class = LlavaNextConfig base_model_prefix = "" supports_gradient_checkpointing = True - _no_split_modules = ["LlamaDecoderLayer"] + _no_split_modules = ["LlamaDecoderLayer", "CLIPEncoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_cache_class = True _supports_flash_attn_2 = True diff --git a/src/transformers/models/llava_onevision/modeling_llava_onevision.py b/src/transformers/models/llava_onevision/modeling_llava_onevision.py index 18475d381a01..1f2766546d05 100644 --- a/src/transformers/models/llava_onevision/modeling_llava_onevision.py +++ b/src/transformers/models/llava_onevision/modeling_llava_onevision.py @@ -43,7 +43,8 @@ ) from ..auto import AutoModel from .configuration_llava_onevision import LlavaOnevisionConfig - +from ..siglip.modeling_siglip import SiglipEncoderLayer +from ..qwen2.modeling_qwen2 import Qwen2DecoderLayer logger = logging.get_logger(__name__) @@ -301,7 +302,7 @@ class LlavaOnevisionPreTrainedModel(PreTrainedModel): config_class = LlavaOnevisionConfig base_model_prefix = "" supports_gradient_checkpointing = True - _no_split_modules = ["LlamaDecoderLayer"] + _no_split_modules = ["Qwen2DecoderLayer", "SiglipEncoderLayer"] _skip_keys_device_placement = "past_key_values" _supports_cache_class = True _supports_flash_attn_2 = True