fix accelerate tests for roberta xl

[SunMarc]

What does this do ?

This PR fixes the no_split_modules for roberta xl model introduced in this PR #31223. Also, we need to change model_split_percents, so that we made sure to have enough space on the first gpu for the embeddings. The tests are passing:

tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py::XLMRobertaXLModelTest::test_cpu_offload If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
[2024-06-06 14:56:19,427] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
PASSED
tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py::XLMRobertaXLModelTest::test_disk_offload_bin If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
PASSED
tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py::XLMRobertaXLModelTest::test_disk_offload_safetensors If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
PASSED
tests/models/xlm_roberta_xl/test_modeling_xlm_roberta_xl.py::XLMRobertaXLModelTest::test_model_parallelism If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`
PASSED

cc @ydshieh since you pinged me on this one

[ydshieh]

if you can leave a comment (in the PR page is ok) why this fix the issue, that would be very nice as knowledge sharing ❤️ .

[SunMarc]

This is needed because there are a lot of residual connections at the level of the layer. So we need to make sure that all the ops that are happening inside of the layer stay on the same device. Otherwise, we would have to do manually move the tensors(.to(device) and this is not good. If you check other _no_split_modules with transformers architecture (llama, mistral ...), you will see that most of them have the layer class inside of it.

[ydshieh]

Thanks a lot!

[ydshieh]

Thank you for fixing it so quickly!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[amyeroberts]

Thanks for fixing!