Add a use_parallel_residual argument to control the residual computing way
Author: NinedayWang
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 -->

Add a gpt_j_residual argument to control the residual computing way, the default value is False, that is 
consistent with https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py#L592. And we can convert the model trained by [gpt-neox](https://github.com/EleutherAI/gpt-neox) into huggingface more easily.



## Who can review?
Anyone in the community is free to review the PR once the tests have passed.
@LysandreJik @patrickvonplaten 

<!-- 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:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->


Comments:
_The documentation is not available anymore as the PR was closed or merged._

Thanks @NinedayWang !

If someone could review this, that would be great. 

This PR will also allow loading our PolyCoder model in `transformers` (https://arxiv.org/pdf/2202.13169.pdf)

Thanks a lot for the PR @NinedayWang,

I'm however not 100% sure we want that as we don't try to make Transformer models be very configurable generally. 
Are there already pretrained checkpoints that would be useful for the community with `gpt_j_residual=True`?

Hi @patrickvonplaten! I appreciate your perspective, but I think in this case supporting the variation is warranted. The default of nearly all [training configurations in the NeoX toolkit](https://github.com/EleutherAI/gpt-neox/tree/main/configs) is to have this flag set to `False`. Only the 20B configuration uses that residual. So I expect that supporting this variation will make deploying new models trained with the NeoX toolkit easier for a lot of folks. Given how costly these models are to train, we are not planning to create a new variant using this residual.

Sorry I don't fully follow here:

- The GPT-NeoX checkpoint can be loaded with this architecture: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py
- So supporting this flag is only for GPT-NeoX, but we already have this checkpoint in Transfomers: https://huggingface.co/EleutherAI/gpt-neox-20b

Ah, yes let me clarify. [GPT-NeoX](https://github.com/EleutherAI/gpt-neox/) is a *toolkit* that can be (and is actively) used to train GPT-style models. It supports a broad range of model sizes, and has a few other hyper-parameters to vary the architecture in other ways, like that `gpt_j_residual` flag.

Now **Neox-20B** is a specific, 20B parameter *model* trained with this toolkit. It largely uses the same configuration that other models trained with GPT-NeoX would, with the notable exception of the aforementioned residual flag: that flag is set to `False` in all configurations by default, but was turned to `True` by the authors of the 20B model. As such, other models trained with the GPT-NeoX toolkit are unlikely to have this flag enabled.

So for HuggingFace/transformers to support most other models trained with the NeoX toolkit, including [PolyCoder](https://github.com/VHellendoorn/Code-LMs/), we could either add multiple other `modeling_gpt_neox_MODELNAME.py*` style architectures, or make the basic `modeling_gpt_neox.py` architecture a bit more flexible. The latter seems more reasonable to me, but if the HF community prefers the former, that could work for us too.

Hope this clarifies things!
-Vincent

> Ah, yes let me clarify. [GPT-NeoX](https://github.com/EleutherAI/gpt-neox/) is a _toolkit_ that can be (and is actively) used to train GPT-style models. It supports a broad range of model sizes, and has a few other hyper-parameters to vary the architecture in other ways, like that `gpt_j_residual` flag.
> 
> Now **Neox-20B** is a specific, 20B parameter _model_ trained with this toolkit. It largely uses the same configuration that other models trained with GPT-NeoX would, with the notable exception of the aforementioned residual flag: that flag is set to `False` in all configurations by default, but was turned to `True` by the authors of the 20B model. As such, other models trained with the GPT-NeoX toolkit are unlikely to have this flag enabled.
> 
> So for HuggingFace/transformers to support most other models trained with the NeoX toolkit, including [PolyCoder](https://github.com/VHellendoorn/Code-LMs/), we could either add multiple other `modeling_gpt_neox_MODELNAME.py*` style architectures, or make the basic `modeling_gpt_neox.py` architecture a bit more flexible. The latter seems more reasonable to me, but if the HF community prefers the former, that could work for us too.
> 
> Hope this clarifies things! -Vincent

Hey @VHellendoorn,

Thanks for clarifying! Putting @LysandreJik and @sgugger in cc here. Given the "single-file" policy of Transformers (see post [here](https://huggingface.co/blog/transformers-design-philosophy)), I think we would indeed prefer to add a new file such as `modeling_poly_coder.py` if the architecture is too different to existing architectures, such as gpt_j or gpt_neox_20b. 
Also just one more question, if PolyCoder follows the same architecture than GPT-NeoX, couldn't we just load it with https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt_neox/modeling_gpt_neox.py ? 

We're definitely more than happy though to help get Polycoder added to Transformers (cc @lvwerra as well)

Hi @patrickvonplaten, 

Thanks, yes that would work for us too. The reason we can't load PolyCoder with that architecture file is precisely because `modeling_gpt_neox.py` hard-codes the assumption that `gpt_j_residual` is set to `True`. Hence the change in this PR, which makes that a configurable boolean. If we add a special `modeling_polycoder.py` file, it will just be identical to the `modeling_gpt_neox.py` one *except* for using the "normal" residual branch that most other models trained with GPT-NeoX will tend to use. So a slightly weird consequence of splitting the architectures across two files would be that most new models trained with GPT-NeoX will have to be loaded with the polycoder architecture, instead of the neox one. This PR would avoid such duplication by making that a togglable boolean instead.

-Vincent


As @patrickvonplaten mentioned, Transformers is not a modular toolkit. It's therefore not surprising that one toolkit class such as GPT-Neo-X in EleutherAI is split in several different classes in Transformers (exactly like BART from fairseq is split in multiple classes here).

Thanks for your reply @patrickvonplaten @sgugger.

Let me do some explanation. [GPT-NeoX](https://github.com/EleutherAI/gpt-neox/) supports two different residual computing ways using `gpt_j_residual` configuration:
https://github.com/EleutherAI/gpt-neox/blob/main/megatron/model/transformer.py#L592
And the default value is `gpt_j_residual=False`:
https://github.com/EleutherAI/gpt-neox/blob/main/megatron/neox_arguments/neox_args.py#L311
```
gpt_j_residual: bool = False
"""
If false, we use the conventional residual path:
  x = x + attn(ln1(x))
  x = x + mlp(ln2(x))
Otherwise, we use the residual path from GPT-J, which offers a slight speedup:
  x = ln(x)
  x = x + attn(x) + mlp(x)
"""
```

As @VHellendoorn said, [gpt-neox-20b](https://huggingface.co/EleutherAI/gpt-neox-20b) is a special case of specifying `gpt_j_residual=True` in the [20B.yml](https://github.com/EleutherAI/gpt-neox/blob/main/configs/20B.yml#L29) config file, most models trained with GPT-NeoX use the default value of False, such as [small.yml](https://github.com/EleutherAI/gpt-neox/blob/main/configs/small.yml), [medium.yml](https://github.com/EleutherAI/gpt-neox/blob/main/configs/medium.yml), [large.yml](https://github.com/EleutherAI/gpt-neox/blob/main/configs/large.yml), [2-7B.yml](https://github.com/EleutherAI/gpt-neox/blob/main/configs/2-7B.yml) and so on. However, `modeling_gpt_neox.py` is an implementation that assumes `gpt_j_residual=True`, so we cannot use `modeling_gpt_neox.py` to load PolyCoder, even though PolyCoder actually follows the same architecture as GPT-NeoX.

I've read about the "single-file" policy, but I think GPT-NeoX is a bit special. If we load `gpt-neox-20b` with `model_type=gpt_neox`, but `gpt-neox-2.7b` or `gpt-neox-0.4b` with `model_type=polycoder`, it can be confusing and people need more time to figure out which model_type is suitable.

Hey @NinedayWang,

Thanks for the explanation. Sorry some more questions to clarify: Why is it called `gpt_j_residual`? Could this be changed to another name? I don't fully understand the relation to GPT-J here.

If we have half the gpt-neox checkpoints using one residual architecture and gpt-neox-20b another architecture I'm actually not too against trying to fit it in one file.

> Hey @NinedayWang,
> 
> Thanks for the explanation. Sorry some more questions to clarify: Why is it called `gpt_j_residual`? Could this be changed to another name? I don't fully understand the relation to GPT-J here.
> 
> If we have half the gpt-neox checkpoints using one residual architecture and gpt-neox-20b another architecture I'm actually not too against trying to fit it in one file.

Thanks a lot! The name `gpt_j_residual` comes from the developers of GPT-NeoX. Actually, the unconventional residual architecture in GPT-NeoX is inherited from [GPT-J](https://github.com/huggingface/transformers/blob/4157e3cd7e2bb5a7be6dc065a3e20c49cc1300ab/src/transformers/models/gptj/modeling_gptj.py#L322). For clarity and to be the same as the original GPT-NeoX, I think it is better to keep the name `gpt_j_residual`.

> > Hey @NinedayWang,
> > Thanks for the explanation. Sorry some more questions to clarify: Why is it called `gpt_j_residual`? Could this be changed to another name? I don't fully understand the relation to GPT-J here.
> > If we have half the gpt-neox checkpoints using one residual architecture and gpt-neox-20b another architecture I'm actually not too against trying to fit it in one file.
> 
> Thanks a lot! The name `gpt_j_residual` comes from the developers of GPT-NeoX. Actually, the unconventional residual architecture in GPT-NeoX is inherited from [GPT-J](https://github.com/huggingface/transformers/blob/4157e3cd7e2bb5a7be6dc065a3e20c49cc1300ab/src/transformers/models/gptj/modeling_gptj.py#L322). For clarity and to be the same as the original GPT-NeoX, I think it is better to keep the name `gpt_j_residual`.

Is this essentially the "parallel" residual computation that allows the model to be tensor-parallelized better (especially for TPUs) - e.g. the same architecture that was used in PALM: https://arxiv.org/abs/2204.02311 ?

> > > Hey @NinedayWang,
> > > Thanks for the explanation. Sorry some more questions to clarify: Why is it called `gpt_j_residual`? Could this be changed to another name? I don't fully understand the relation to GPT-J here.
> > > If we have half the gpt-neox checkpoints using one residual architecture and gpt-neox-20b another architecture I'm actually not too against trying to fit it in one file.
> > 
> > 
> > Thanks a lot! The name `gpt_j_residual` comes from the developers of GPT-NeoX. Actually, the unconventional residual architecture in GPT-NeoX is inherited from [GPT-J](https://github.com/huggingface/transformers/blob/4157e3cd7e2bb5a7be6dc065a3e20c49cc1300ab/src/transformers/models/gptj/modeling_gptj.py#L322). For clarity and to be the same as the original GPT-NeoX, I think it is better to keep the name `gpt_j_residual`.
> 
> Is this essentially the "parallel" residual computation that allows the model to be tensor-parallelized better (especially for TPUs) - e.g. the same architecture that was used in PALM: https://arxiv.org/abs/2204.02311 ?

Yes, it's the same "parallel" architecture as PALM, which provides faster training speed when training large-scale models.