Enable multi-device for more models (#30379)

* feat: support for vitmatte

* feat: support for vivit

* feat: support for beit

* feat: support for blip :D

* feat: support for data2vec

src/transformers/models/beit/modeling_beit.py

@@ -563,6 +563,7 @@ class BeitPreTrainedModel(PreTrainedModel):
     base_model_prefix = "beit"
     main_input_name = "pixel_values"
     supports_gradient_checkpointing = True
+    _no_split_modules = ["BeitLayer"]
 
     def _init_weights(self, module):
         """Initialize the weights"""

src/transformers/models/blip/modeling_blip_text.py

@@ -549,6 +549,7 @@ class BlipTextPreTrainedModel(PreTrainedModel):
 
     config_class = BlipTextConfig
     base_model_prefix = "bert"
+    _no_split_modules = []
 
     def _init_weights(self, module):
         """Initialize the weights"""

src/transformers/models/data2vec/modeling_data2vec_vision.py

@@ -574,6 +574,7 @@ class Data2VecVisionPreTrainedModel(PreTrainedModel):
     base_model_prefix = "data2vec_vision"
     main_input_name = "pixel_values"
     supports_gradient_checkpointing = True
+    _no_split_modules = ["Data2VecVisionLayer"]
 
     def _init_weights(self, module):
         """Initialize the weights"""

src/transformers/models/vitmatte/modeling_vitmatte.py

@@ -73,6 +73,7 @@ class VitMattePreTrainedModel(PreTrainedModel):
     config_class = VitMatteConfig
     main_input_name = "pixel_values"
     supports_gradient_checkpointing = True
+    _no_split_modules = []
 
     def _init_weights(self, module):
         if isinstance(module, nn.Conv2d):

src/transformers/models/vivit/modeling_vivit.py

@@ -387,6 +387,7 @@ class VivitPreTrainedModel(PreTrainedModel):
     base_model_prefix = "vivit"
     main_input_name = "pixel_values"
     supports_gradient_checkpointing = True
+    _no_split_modules = []
 
     def _init_weights(self, module):
         """Initialize the weights"""