Update no split modules in T5Gemma model (#40810)

* Update no split modules in T5Gemma model

* Update no_split_modules also for T5Gemma modular

* Remove model_split_percents from test cases

---------

Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>

Author: npuichigo

src/transformers/models/t5gemma/modeling_t5gemma.py

@@ -585,7 +585,7 @@ class T5GemmaPreTrainedModel(PreTrainedModel):
     config: T5GemmaConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
-    _no_split_modules = ["T5GemmaBlock"]
+    _no_split_modules = ["T5GemmaEncoderLayer", "T5GemmaDecoderLayer"]
     _skip_keys_device_placement = ["past_key_values"]
     _supports_flash_attn = True
     _supports_sdpa = True

src/transformers/models/t5gemma/modular_t5gemma.py

@@ -476,7 +476,7 @@ class T5GemmaPreTrainedModel(Gemma2PreTrainedModel):
     config: T5GemmaConfig
     base_model_prefix = "model"
     supports_gradient_checkpointing = True
-    _no_split_modules = ["T5GemmaBlock"]
+    _no_split_modules = ["T5GemmaEncoderLayer", "T5GemmaDecoderLayer"]
 
     def _init_weights(self, module):
         # TODO: support initialization for encoders and decoders separately(?)

tests/models/t5gemma/test_modeling_t5gemma.py

@@ -597,7 +597,6 @@ class T5GemmaModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTesterMi
     test_pruning = False
     _is_stateful = True
     is_encoder_decoder = True
-    model_split_percents = [0.5, 0.6]
 
     # used in `test_torch_compile_for_training`
     _torch_compile_train_cls = T5GemmaForConditionalGeneration if is_torch_available() else None
@@ -1460,7 +1459,6 @@ class T5GemmaEncoderOnlyModelTest(ModelTesterMixin, unittest.TestCase):
     test_headmasking = False
     _is_stateful = True
     is_encoder_decoder = False
-    model_split_percents = [0.4, 0.5]
 
     # won't fix
     test_torchscript = False