Fix `test_model_parallelism` (#25359)

* fix

* fix

---------
Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>

tests/models/roberta_prelayernorm/test_modeling_roberta_prelayernorm.py

@@ -395,6 +395,7 @@ class RobertaPreLayerNormModelTest(ModelTesterMixin, GenerationTesterMixin, Pipe
         else {}
     )
     fx_compatible = False
+    model_split_percents = [0.5, 0.8, 0.9]
 
     def setUp(self):
         self.model_tester = RobertaPreLayerNormModelTester(self)

tests/models/vilt/test_modeling_vilt.py

@@ -235,6 +235,7 @@ class ViltModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCase):
     test_pruning = False
     test_headmasking = False
     test_torchscript = False
+    model_split_percents = [0.5, 0.8, 0.9]
 
     # ViltForMaskedLM, ViltForQuestionAnswering and ViltForImagesAndTextClassification require special treatment
     def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):

tests/models/vit_hybrid/test_modeling_vit_hybrid.py

@@ -163,6 +163,7 @@ class ViTHybridModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCas
     test_pruning = False
     test_resize_embeddings = False
     test_head_masking = False
+    model_split_percents = [0.5, 0.9]
 
     def setUp(self):
         self.model_tester = ViTHybridModelTester(self)

tests/models/xglm/test_modeling_xglm.py

@@ -347,6 +347,10 @@ class XGLMModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMixin
             model = XGLMModel.from_pretrained(model_name)
             self.assertIsNotNone(model)
 
+    @unittest.skip("Does not work on the tiny model as we keep hitting edge cases.")
+    def test_model_parallelism(self):
+        super().test_model_parallelism()
+
 
 @require_torch
 class XGLMModelLanguageGenerationTest(unittest.TestCase):

tests/test_modeling_common.py

@@ -2597,7 +2597,7 @@ class ModelTesterMixin:
 
             model_size = compute_module_sizes(model)[""]
             # We test several splits of sizes to make sure it works.
-            max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents]
+            max_gpu_sizes = [int(p * model_size) for p in self.model_split_percents[1:]]
             with tempfile.TemporaryDirectory() as tmp_dir:
                 model.cpu().save_pretrained(tmp_dir)