From 99895ae5e265ea67be77d9bc51544b093105424c Mon Sep 17 00:00:00 2001
From: Marc Sun <57196510+SunMarc@users.noreply.github.com>
Date: Thu, 6 Jun 2024 15:44:35 +0200
Subject: [PATCH] fix accelerate tests for roberta xl (#31288)

* fix accelerate tests for roberta xl

* style
---
 .../models/xlm_roberta_xl/modeling_xlm_roberta_xl.py            | 2 +-
 tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py     | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/transformers/models/xlm_roberta_xl/modeling_xlm_roberta_xl.py b/src/transformers/models/xlm_roberta_xl/modeling_xlm_roberta_xl.py
index 7534eb4dad3b4c..d189086cd89e69 100644
--- a/src/transformers/models/xlm_roberta_xl/modeling_xlm_roberta_xl.py
+++ b/src/transformers/models/xlm_roberta_xl/modeling_xlm_roberta_xl.py
@@ -572,7 +572,7 @@ class XLMRobertaXLPreTrainedModel(PreTrainedModel):
 
     config_class = XLMRobertaXLConfig
     base_model_prefix = "roberta"
-    _no_split_modules = ["XLMRobertaXLEmbeddings", "XLMRobertaXLSelfAttention"]
+    _no_split_modules = ["XLMRobertaXLEmbeddings", "XLMRobertaXLLayer"]
 
     # Copied from transformers.models.bert.modeling_bert.BertPreTrainedModel._init_weights
     def _init_weights(self, module):
diff --git a/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py b/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py
index d8a37d47790945..22663db27c811c 100644
--- a/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py
+++ b/tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py
@@ -387,6 +387,8 @@ class XLMRobertaXLModelTest(ModelTesterMixin, GenerationTesterMixin, PipelineTes
         else {}
     )
 
+    model_split_percents = [0.5, 0.85, 0.95]
+
     # TODO: Fix the failed tests
     def is_pipeline_test_to_skip(
         self, pipeline_test_casse_name, config_class, model_architecture, tokenizer_name, processor_name