Repository: huggingface/transformers
Issue category: Feature Request

Issue:
SiqlipVisionModel does not support "device map= auto": no split modules`attribute
### Feature request

How to make SiglipVisionModel can support Auto map to map to multiple GPUs.

### Motivation

Currently, using MLLM with siglip, the whole model might need automap, since the vision encoder part is SiglipVisionModel, if it doesn;t support , would be a big problme

Also, would consider add flashattention support for Siglip?

### Your contribution

None for now


Hints:
Hi @lucasjinreal, thanks for opening a feature request! 

Could you share a code snippet of how the model is being created with `auto_map` and the running environment (run `transformers-cli env` in the terminal and copy-paste the output)? SigLip should support `device_map="auto"`
I checked with the code in [model-doc](https://huggingface.co/docs/transformers/en/model_doc/llava) for SigLip and also got an error. Supporting "device_map" in VLMs is indeed needed important. I believe `_no_split_modules` should be same as in [ClipModel](https://github.com/huggingface/transformers/blob/main/src/transformers/models/clip/modeling_clip.py)

For `flashattention` afaik current VLMs in transformers use optimized attn implementations only for LLM backbone (e.g. LLaVa supports Flash-Attn and SDPA even though CLIP doesn't). There's an issue for adding SDPA attn (#30565) to all VLMs, I can open another tracker-issue for Flash-Attn but will not able to work on it right now. Open for community contributions 
I have surpassed this error, by simply add a _no_split_modules = [] to the attribute.

But it could be better add inside transoformers, it's just a single line. I could submit a PR for this.

As for flashattn, it's a really needs, it can boost vlms training more faster.
@lucasjinreal cool, PR would be nice but you need to test in multi-gpu setting that everything is being split correctly. I don't think that an empty "split_modules" will work as the most similar CLIP doesn't split at some modules. If you don't have multiple gpus, I can run some tests after the PR is open :)

Flash-Attn noted, thanks, will add to my todo list!
I have been using empty list and make it can be deivicemap auto on multiple GPUs, currently inference is normal. I still didn't know why CLIPVIsionModel should make CLIPENcoderLayer didn't automap though.
@lucasjinreal i just noticed that SigLip already has `_no_split_modules` in TextModel and in VisionModel, yet not in the SigLipModel. If I do `_no_split_modules=[]` as you tried, device mismatch error is raised so we have to add text and vision models' `_no_split_modules` to enable it

LMK if you're up to opening a PR :)
Hi, In my cased I just using SIglipVisionModel as a parent class and used a SiglipVisionModelSplit(SiglipVisionModel) in my MLLM.

So I think it not appliable to inside of transformers.  Let me think a better way to do this


I believe the best solution is to copy 'no-split-modules' that are already indicated in text-vision components, and add them in SiglipModel's 'no-split-modules'

Patch:
diff --git a/src/transformers/models/siglip/modeling_siglip.py b/src/transformers/models/siglip/modeling_siglip.py
index d605f49261ae6f..4c534bbce6ce8a 100644
--- a/src/transformers/models/siglip/modeling_siglip.py
+++ b/src/transformers/models/siglip/modeling_siglip.py
@@ -496,6 +496,13 @@ class SiglipPreTrainedModel(PreTrainedModel):
     config_class = SiglipConfig
     base_model_prefix = "siglip"
     supports_gradient_checkpointing = True
+    _no_split_modules = [
+        "SiglipTextEmbeddings",
+        "SiglipEncoderLayer",
+        "SiglipVisionEmbeddings",
+        "SiglipEncoderLayer",
+        "SiglipMultiheadAttentionPoolingHead",
+    ]
 
     def _init_weights(self, module):
         """Initialize the weights"""
@@ -816,8 +823,6 @@ def forward(
 class SiglipTextModel(SiglipPreTrainedModel):
     config_class = SiglipTextConfig
 
-    _no_split_modules = ["SiglipTextEmbeddings", "SiglipEncoderLayer"]
-
     def __init__(self, config: SiglipTextConfig):
         super().__init__(config)
         self.text_model = SiglipTextTransformer(config)
@@ -959,7 +964,6 @@ def forward(self, hidden_state):
 class SiglipVisionModel(SiglipPreTrainedModel):
     config_class = SiglipVisionConfig
     main_input_name = "pixel_values"
-    _no_split_modules = ["SiglipVisionEmbeddings", "SiglipEncoderLayer", "SiglipMultiheadAttentionPoolingHead"]
 
     def __init__(self, config: SiglipVisionConfig):
         super().__init__(config)
@@ -1222,7 +1226,10 @@ def forward(
         text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
 
         # cosine similarity as logits
-        logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * self.logit_scale.exp() + self.logit_bias
+        logits_per_text = (
+            torch.matmul(text_embeds, image_embeds.t().to(text_embeds.device)) * self.logit_scale.exp()
+            + self.logit_bias
+        )
         logits_per_image = logits_per_text.t()
 
         loss = None


Test patch:
diff --git a/tests/models/siglip/test_modeling_siglip.py b/tests/models/siglip/test_modeling_siglip.py
index 12ac11251dc4c9..af5d0bf2bc3e83 100644
--- a/tests/models/siglip/test_modeling_siglip.py
+++ b/tests/models/siglip/test_modeling_siglip.py
@@ -443,6 +443,12 @@ class SiglipModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
     test_pruning = False
     test_resize_embeddings = False
     test_attention_outputs = False
+    # MP works but offload doesn't work when the MultiheadAttention is offloaded
+    # TODO: One potential solution would be to add to set preload_module_classes = ["SiglipMultiheadAttentionPoolingHead"]
+    # in the dispatch_model function
+    test_cpu_offload = False
+    test_disk_offload_safetensors = False
+    test_disk_offload_bin = False
 
     # Copied from tests.models.clip.test_modeling_clip.CLIPModelTest.setUp with CLIP->Siglip
     def setUp(self):
@@ -618,6 +624,12 @@ class SiglipForImageClassificationModelTest(ModelTesterMixin, PipelineTesterMixi
     test_pruning = False
     test_resize_embeddings = False
     test_attention_outputs = False
+    # MP works but offload doesn't work when the MultiheadAttention is offloaded
+    # TODO: One potential solution would be to add to set preload_module_classes = ["SiglipMultiheadAttentionPoolingHead"]
+    # in the dispatch_model function
+    test_cpu_offload = False
+    test_disk_offload_safetensors = False
+    test_disk_offload_bin = False
 
     def setUp(self):
         self.model_tester = SiglipForImageClassificationModelTester(self)