Add GPT-J 6B support to the gpt-neo implementation
Author: finetunej
Pull request: yes

# What does this PR do?

This PR mainly adds support for the GPT-J 6B model. A [conversion script](https://gist.github.com/finetuneanon/ee196c6cd16af1de4ca444862414683a) and [config.json](https://gist.github.com/finetuneanon/a55bdb3f5881e361faef0e96e1d41f09) for the slim checkpoint are also available.

It also addresses the local attention issue from #11320 in the same way as PR #11630 and works around an issue with torch.multinomial that will allow zero probability tokens to be chosen when sampling from an fp16 model.

Special thanks to the great folks of the EleutherAI discord, who helped me debug the RoPE implementation and to @kurumuz (NovelAI) who worked on this as well.

<!--
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 -->

Fixes #11320


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/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/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/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.

<!-- 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:
Just as a note, we have a PyTorch checkpoint for GPT-J that we will be ready to upload once this PR goes through.

The conversion script in the top post generates a single file checkpoint, but for models of this size, I've found split up checkpoints usually more efficient to load and handle. Such split up checkpoints can be generated using [this conversion script](https://gist.github.com/finetuneanon/7dd417a31338a63f219a49702e0550db) and loaded as follows:

```python
try:
    from collections.abc import MutableMapping
except ImportError:
    from collections import MutableMapping
from pathlib import Path

class Checkpoint(MutableMapping):
    def __init__(self, chkpt_dir, device="cpu"):
        self.device = device
        self.chkpt_dir = Path(chkpt_dir)
        self.checkpoint = torch.load(str(chkpt_dir / Path("m.pt")))
    def __len__(self):
        return len(self.checkpoint)
    def __getitem__(self, key):
        path = self.chkpt_dir / Path(self.checkpoint[key]).name
        return torch.load(str(path), map_location=self.device)
    def __setitem__(self, key, value):
        return
    def __delitem__(self, key, value):
        return
    def keys(self):
        return self.checkpoint.keys()
    def __iter__(self):
        for key in self.checkpoint:
            yield (key, self.__getitem__(key))
    def __copy__(self):
        return Checkpoint(self.chkpt_dir, device=self.device)
    def copy(self):
        return Checkpoint(self.chkpt_dir, device=self.device)

from transformers import GPTNeoForCausalLM, AutoConfig
config = AutoConfig.from_pretrained(model_name)
model = GPTNeoForCausalLM.from_pretrained(pretrained_model_name_or_path=None, config=config, state_dict=Checkpoint("checkpoint"))
```

Having a more integrated or better specifid way of loading them would be helpful, but I'm not sure what the best place for that would be.

**Edit: Updated to handle renamed checkpoint folders.**

I noticed that there was a typo in the config file linked from the PR text, which caused it to be invalid JSON. It's fixed now.

Also ran some evaluations using the [eval harness](https://github.com/EleutherAI/lm-evaluation-harness) on the ported model now:

|   Task   |    Metric     |Value |
|----------|---------------|-----:|
|lambada   |ppl            |4.1060|
|          |ppl_stderr     |0.0886|
|          |acc            |0.6833|
|          |acc_stderr     |0.0065|
|winogrande|acc            |0.6480|
|          |acc_stderr     |0.0134|
|piqa      |acc            |0.7541|
|          |acc_stderr     |0.0100|
|          |acc_norm       |0.7612|
|          |acc_norm_stderr|0.0099|
|hellaswag |acc            |0.4895|
|          |acc_stderr     |0.0050|
|          |acc_norm       |0.6614|
|          |acc_norm_stderr|0.0047|

The eval numbers are a little shy of what we have for the Jax model, but close enough that FP rounding could plausibly explain the difference:

Lambada: 3.99 ppl, 0.697 acc
Winogrande: 0.653
PiQA: 0.765
HellaSwag: 0.661

It should also be noted that my results were with fp16. It should be easy enough to modify the conversion script to cast to fp32 (just replace `half()` with `float()`), which might give results closer to the original evaluation, but I don't currently have access to hardware which could run the model in fp32 at a reasonable speed to do an evaluation on it.

> It should also be noted that my results were with fp16. It should be easy enough to modify the conversion script to cast to fp32 (just replace `half()` with `float()`), which might give results closer to the original evaluation, but I don't currently have access to hardwhere which could run the model in fp32 at a reasonable speed to do an evaluation on it.

Oh yes, I concur. That wasn’t meant as a detraction at all. I’m not sure if EAI had enough free GPUs, but I can look at getting the evals run at full precision later this week.

Just curious - how long before new models are merged to the repo, generally speaking? And how long until it's available in the hosted inference API?

Hi @finetuneanon 

Amazing, thanks a lot for porting `GPT-J` so quickly!

The changes to local attention look good to me. But would be nice to split the PR into two 
1. simplify local attention
2. and add GPT-J in a new model file.

While supporting `GPT-J` in the `GPTNeo` model is totally doable, we would like to avoid that. The overall philosophy is to combine more than one model only if the forward pass is exactly similar or requires some really minor changes. If I understand it correctly, here are the differences between `GPT-J` and `GPTNeo` :
- `GPT-J` uses rotary embeddings
- It scales attention weights
- no bias in the attention projection layer (the `out_proj` layer in attention class)
- does not use layer_norm before the feed forward layer (`mlp`)
- no residual connection between `hidden_states` and `attention_output` , just one residual connection which is added to `attention + mlp(hiddn)`
- uses bias in the output projection layer
- does not ties word embeddings with output layer

The current PR handles this using the `config.jax` argument, but it’s a bit confusing, and generally, one config param should not control these many changes. So if we really decide to support this in `GPTNeo` we would probably end-up with different config params like `attention_projection_bias`, `output_bias`, `attention_residual`,  `scale_attention`. So it’s cleaner IMO to add a new model for this.

Also, `Transformers` isn’t really a modular toolkit for building different models, The goal is to keep every model file responsible for one model so it becomes easier for everyone to understand it and modify it according to their needs.  And also makes it easy for us to maintain these different models.
cc @LysandreJik , @sgugger , @patrickvonplaten

Happy to help in any way to add the model :) 

To be quite honest, I think reading a couple of if branches makes the differences between models much clearer than having to compare two completely different model classes with small differences. You mention that transformers is not intended to be a modular framework, so there should be no issue with controlling these changes through a single configuration variable, although perhaps the naming of `jax` is not optimal. I would be open to changing this to e.g. `gptj`. Splitting it up into multiple options would only make sense to actually turn it into a modular framework.

I would also prefer not splitting the pull request.

Let's agree to disagree: this is one of the core principle of the Transformers library, explicitly stated in our [philosophy](https://huggingface.co/transformers/philosophy.html). We've been proceeding like this since the beginning, and while we definitely understand where you're coming from, this is a defining principle of our library which we are not eager to change as it has been validated both by [surveys](https://discuss.huggingface.co/t/transformers-huge-community-feedback/120) and by community feedback.

Unfortunately, we will have to insist on GPT-J following the same approach as the rest of the models - for philosophy, maintenance and coherence's sake. Let us know if you would like us to take over, we are happy to! Thank you for your understanding.

Yes, please take over in that case.

> Let's agree to disagree: this is one of the core principle of the Transformers library, explicitly stated in our [philosophy](https://huggingface.co/transformers/philosophy.html). We've been proceeding like this since the beginning, and while we definitely understand where you're coming from, this is a defining principle of our library which we are not eager to change as it has been validated both by [surveys](https://discuss.huggingface.co/t/transformers-huge-community-feedback/120) and by community feedback.
> 
> Unfortunately, we will have to insist on GPT-J following the same approach as the rest of the models - for philosophy, maintenance and coherence's sake. Let us know if you would like us to take over, we are happy to! Thank you for your understanding.

If the current changes were to be submitted as a new model, instead of a modification to GPT-Neo, would there be any significant further changes to be made?

> If the current changes were to be submitted as a new model, instead of a modification to GPT-Neo, would there be any significant further changes to be made?

The new modeling file is the main thing, we have a [template to add new models](https://github.com/huggingface/transformers/tree/master/templates/adding_a_new_model) that should take care of everything else.

I have opened a PR that attempts to refit this into HF's paradigm. I recommend closing this PR and directing discussion to #12243 12243

As discussed above, this PR will be split into two
- GPT-J (which Stella has already started)
- simplifying GPTNeo local attention 

Closing this PR now.