Fix beam search when using model parallel
Author: pfldy2850
Pull request: yes

# What does this PR do?

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

This PR fixes a crash when running beam search on multiple GPUs. Similar issue is also observed and fixed on T5 https://github.com/huggingface/transformers/pull/11717 and LLama https://github.com/huggingface/transformers/pull/24224

Fixes # (issue)


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.
@ArthurZucker @younesbelkada 
<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker and @younesbelkada
- vision models: @amyeroberts
- speech models: @sanchit-gandhi
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @gante
- pipelines: @Narsil
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @sgugger

Integrations:

- deepspeed: HF Trainer/Accelerate: @pacman100
- ray/raytune: @richardliaw, @amogkam

Documentation: @sgugger, @stevhliu and @MKhalusova

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: @sgugger
- TensorFlow: @Rocketknight1

 -->


Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_24969). All of your documentation changes will be reflected on that endpoint.

I agree with you!
Which do you think is better, fixing the rest of the model in this PR or creating a new PR fixing the rest?

I am willing to do further work based on your comments.

I think it will be easier to have everything in one PR, given how small and repetitive of a change it is! 

@ArthurZucker 

I have added fixed commits for every model that required correction, and I have also made modifications to the cookiecutter template.

And I have updated the PR title and content to align with the task.

As there were numerous models that needed correction, there is a possibility that some parts might have been overlooked. Therefore, I would appreciate it if you could review the changes again. Thank you for your attention to this matter.

There is no `device` in tensorflow, let's limit the changes to pytorch! 

Oh! Thank you for correcting the mistake.
As you suggested, I have dropped the modifications for the tf base model.

@ArthurZucker 

Is there anything else you'd like me to fix?
I want to use the merged main branch for my work.

Not at all sorry, let me have a final look and I'll merge this! 

Hmm, by the way.
It seems like there's already a test in the script you provided to test beam search on multi GPUs.

https://github.com/pfldy2850/transformers/blob/e07126aac6840568b0db0b369d199f3a0cefa28f/tests/test_modeling_common.py#L2468-L2494

Why was this test not conducted for this issue beforehand?

Your are right! Might be the `test_model_parallel`, it is set `False` by default 

@ArthurZucker 

What do you think about setting the `test_model_parallel=True` in the existing modeling test file instead of creating a new test?


Great idea, the problem is that this might also trigger other tests, and there might be a reason why don't test them (maybe too slow / model doesn't need these test as it is not used that much). Pinging @amyeroberts for a final answer 🤗 

OK, I think this is hitting on some areas of our testing suite which are non-obvious and / or need updating. 

AFAICT, there are only two models which have `test_model_parallel=True`, [GPT2 and T5](https://github.com/search?q=repo%3Ahuggingface%2Ftransformers+%22test_model_parallel+%3D+%22&type=code). The tests which check this flag, both use a [deprecated method](https://github.com/huggingface/transformers/blob/5744482abc472e472874b632d23c726affed8650/src/transformers/models/t5/modeling_t5.py#L929) `model.parallelize` - [1](https://github.com/huggingface/transformers/blob/5744482abc472e472874b632d23c726affed8650/tests/test_modeling_common.py#L2413), [2](https://github.com/huggingface/transformers/blob/5744482abc472e472874b632d23c726affed8650/tests/test_modeling_common.py#L2454)- and so this flag is to control testing for backwards compatibility features. 

We have another test which [tests for parallelism](https://github.com/huggingface/transformers/blob/5744482abc472e472874b632d23c726affed8650/tests/test_modeling_common.py#L2584) which doesn't check `test_model_parallel`: test_model_parallelism, which is an accelerate test and checks [if the model has `_no_split_modules` implemented](https://github.com/huggingface/transformers/blob/5744482abc472e472874b632d23c726affed8650/tests/test_modeling_common.py#L2588). 

In addition, generate specific tests should be added to [GenerationTesterMixin](https://github.com/huggingface/transformers/blob/5744482abc472e472874b632d23c726affed8650/tests/generation/test_utils.py#L77), rather than ModelTesterMixin, as only models with `.generate` methods should be tested. 

What I would suggest is: 
* Moving `test_model_parallel_beam_search` to  GenerationTesterMixin. 
* Update the modelings tests so that each of the model's updated in this PR have `all_generative_model_classes` added as attributes to their model tester.
* Update `test_model_parallel_beam_search` to check if the model class has `_no_split_modules` implemented. If not, skip. If it does, then load using `device_map="auto"` instead of `test_model_parallel`. 
* Make sure to mark that the test requires multi gpus


@amyeroberts @ArthurZucker 

I'm sorry, it has been delayed due to busy work. 

I've done the fix as you said.

1. `test_model_parallel_beam_search` function has been moved from ModelTextMixin to GenerationMixin.
2. Based on whether `_no_split_modules` is implemented, skip logic has been written.
3. Parallelized to multiple devices using `device_map="auto"`.
4. Errors caused by the new test have been fixed.

Is there any update?
I need this feature for my academic project.

Friendly ping @ArthurZucker :) 

@dev-cotyledon If you need this for your project, it's possible to work from this branch until it's been merged into main: 

* Clone the repo: `git clone git@github.com:huggingface/transformers.git`
* Create a local environment e.g. `python -m venv my-env`
* Install packages from source in dev mode `cd transformers && pip install -e .`
* Add fork to remote `git remote add pfldy2850 https://github.com/pfldy2850/transformers.git`
* Fetch this branch `git fetch pfldy2850 fix-gpt-neox-parallelize-beam`
* Checkout this branch `git checkout fix-gpt-neox-parallelize-beam`

Your environment will now be running the version of transformers of this branch. 

How kind of you! Thanks a lot sir 🙏

I had to modify the author of the commit, so I amend and then force push.

@ArthurZucker Are you happy for us to merge? 

Yep! Sorry must have missed the ping 😢 

Thanks for the fix @pfldy2850, tests are all green so merged the PR 😉