GPT-J-6B
Author: StellaAthena
Pull request: yes

# What does this PR do?

Introduces the long awaited `GPT J` model class to HuggingFace! Concurrently with this PR being merged I will make a GPT J 6B checkpoint public on the EleutherAI HF page for people to use. The model has been evaluated as being within error tolerances of the GPT J 6B model we released in Jax two months ago.

@patil-suraj was very helpful in assisting me to understand HF philosophy and how to make this PR most in line with the rest of the codebase. Other than that, the major design consideration was to make the configs compatible with GPT-2 rather than GPT-Neo. GPT-Neo has some usability limitations due to its configs having names unrelated to GPT-2’s (see #12183 for details). Given those problems and my hope that GPT-Neo will have it’s configs updated in the future, it seemed like a clear choice to align GPT J with GPT-2.

Shout outs to @finetuneanon whose implementation this one is based off of, as well as @kumuruz for assistence optimizing and debugging.

Supersedes #12243 #13010 #13022 

Closes #12098

## 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?
- [X] 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.
It was discussed in Slack with @patil-suraj 
- [X] 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).
- [X] Did you write any new necessary tests?


## Who can review?

- gpt2: @patrickvonplaten, @LysandreJik, @patil-suraj 

Comments:
There are six failed tests.

Four of them relate to Flax and TF models which I did not add. I may have left some boilerplate code indicating the existence of such models by accident.

One of them relates to docstring issues. I’ll double check the docstrings, but these issues have no impact on the functionality of the model.

One of them appears to be a basic quality assurance check. The code says
```
assert 2 == 3
def test_answer():
>       assert 1 + 1 == 3
E       assert 2 == 3

test_sample.py:2: AssertionError
```
but I have no idea what this means. Assistence it advice would be appreciated.

Looks like a boilerplate test that never got filled out/removed

I don't see any test_sample.py in your branch. I think this file is not committed and probably some example test that somehow ended up in your local copy?

There will also be a tricky merge with the result of #13023, let us know if you need any help with that.

I'm glad y'all approve :)

I made many of the recommended changes, but have some stuff to take care of today. I'll work on addressing them all and tag you guys when it's ready for re-review.

I have been trying to work on this, but I have [a few](https://xkcd.com/1070/) lingering questions that have been slowing my  progress.
1. Is there a certain motivation for [`Attention`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/openai/modeling_openai.py#L156-L157), [`GPT2Attention`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt2/modeling_gpt2.py#L150-L155), [`MLP`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/openai/modeling_openai.py#L237-L238), and [`GPT2MLP`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt2/modeling_gpt2.py#L275-L276) to use [`torch.nn.Conv1D`](https://pytorch.org/docs/1.9.0/generated/torch.nn.Conv1d.html)? [`GPTNeoSelfAttention`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt_neo/modeling_gpt_neo.py#L321-L324), [`GPTNeoLocalSelfAttention`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt_neo/modeling_gpt_neo.py#L390-L393) and [`GPTNeoMLP`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt_neo/modeling_gpt_neo.py#L521-L522) buck the trend by using [`torch.nn.Linear`](https://pytorch.org/docs/1.9.0/generated/torch.nn.Linear.html) instead. The last instance also maintains the inappropriate prefix of `c_` in the layer name (which I presume is to indicate that it is a convolutional layer).
2. [`GPT2Config`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt2/configuration_gpt2.py#L188-L202) and [`OpenAIGPTConfig`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/openai/configuration_openai.py#L162-L176) both alias four arguments of the constructor by defining four new properties. [`GPTNeoConfig`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt_neo/configuration_gpt_neo.py#L174-L180) (in its perpetual inconsistency) only maintains two of them while using names that are also not consistent with `OpenAIGPTConfig` and `GPT2Config`. Why maintain these redundant properties? If the intention is to use them to rename constructor arguments rather than to provide dual access, they could be simply renamed in the initial block of the constructor. Even worse, many values that are initialized from the config do not maintain continuality in their naming, a problem that spans all four models discussed here. (To demonstrate, trace how [`GPT2Config.n_embd` is aliased to `GPT2Config.hidden_size`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt2/configuration_gpt2.py#L192-L194) so that it can be [later accessed via `GPT2Config.hidden_size` and assigned to `GPT2Attention.embed_dim`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gpt2/modeling_gpt2.py#L138))
3. Are we maintaining [`GPTJForSequenceClassification`](https://github.com/huggingface/transformers/blob/24ac25a07ba80b9b7b6e396887305437478398ff/src/transformers/models/gptj/modeling_gptj.py#L781-L892) or not? @StellaAthena [removed the reference to it](https://github.com/huggingface/transformers/pull/13022/commits/4efbbeca8fe4674abf42d253b6dcdd70077cdebf) but @patil-suraj [restored that reference](https://github.com/huggingface/transformers/pull/13022/commits/24ac25a07ba80b9b7b6e396887305437478398ff). Defining the scope of this PR is likely a good idea to prevent counterproductive work.

Hi @EricHallahan 

1. The original GPT2 used the conv1d layer instead of the `linear` (it's essentially linear but just keeps weights transposed). It's rather confusing why they chose that name, so we try not to use it anymore. The names are there just for historical reasons :D, Linear is well known and easy to understand and no transpose is needed when doing the computation.

2. The `hidden_size`, `max_position_embeddings`, `num_hidden_layers`, `num_attention_heads` are common attribute across all configs. They help enable some common tests for configs and models. Note that GPT2 and GPT were added before introducing these attributes so now they are aliased. Simply renaming the constructor argument is not an option since it'll break backward compatibility.  If constructor arguments are renamed thousands of GPT2 models on the hub will fail to load since their config is already defined. As much as we would like to do that, it's not an option.
       Also, note that when GPT2 was added `transformers` was a bit new and has evolved a lot since then so things and guidelines have changed a bit. In general, we try to use these new names whenever possible, but for this model, I think it's fine to use the `n_layer`, `n_head` etc for consistency with the GPT2 config since Leo at some point had mentioned that it's useful to have those to be able to swap models easily. 
     But so far there were no issues from the community about these names, so I fail to see what's the big problem here.

3. ` GPTJForSequenceClassification` was type hinted in the main init, if a model class is used somewhere in main init then the quality tests requires that such class should be tested, added to the auto model, and documented. So the tests were failing, so I decided to add it back since the class was already defined. But no strong opinion about it, feel free to completely remove it.

Hope this answers your question :) 

@patil-suraj That perfectly answers my questions. I had been making the assumption that GPT-2 was the benchmark for how the model should be structured, but relaxing that assumption resolves that conflict.
I've been working up cleaning up the mess left from the Attention mixin and consolidating those classes to the unified `GPTJAttention`. I should get some sleep, but hopefully I'll get that pushed sometime tomorrow.

Glad to know! I agree that the AttentionMixin is rather confusing. Thanks for working on it. Apart from few comments above the PR is already in very good shape! 

I have committed the suggested changes by @sgugger, or at least as many as I could before my phone interface started acting weird. I’m coming back from vacation tomorrow and can go over the PR for real when I get back.

@patil-suraj thanks for the info! This is quite helpful. I have no problem supporting sequence classification if the code supports it. I had removed it because I didn’t think we could support it without writing a bunch more code.

It looks like we are quite close to getting this merged! Thanks for all the help @EricHallahan @kurumuz 

@StellaAthena thanks for your work on this PR! Can you share details on the GPU (or AWS instance type) you used for testing your code? I plan to run some benchmarking with this model after this PR is merged (hopefully soon!).

> @StellaAthena thanks for your work on this PR! Can you share details on the GPU (or AWS instance type) you used for testing your code? I plan to run some benchmarking with this model after this PR is merged (hopefully soon!).

The original model was trained on TPUs, and my testing of this PyTorch port has been on 8x A100 clusters

Okay, well if that's the case, I'd love to see stats backing the claim that
I misunderstand.

Last I checked, fp32 weights for a 2.7B model itself is ~10 GB of disk
space. So this 6B model you are creating a PR for will likely take more
than 16 GB of disk space. So I do not see how you "absolutely can perform
inference on a 16 GB V100 GPU" with this GPT-J 6B model. Please
advise/correct my understanding since I'm still very much a newbie. Thanks!


On Mon, Aug 16, 2021, 9:38 PM Stella Biderman ***@***.***>
wrote:

> ***@***.**** commented on this pull request.
> ------------------------------
>
> In src/transformers/models/gptj/modeling_gptj.py
> <https://github.com/huggingface/transformers/pull/13022#discussion_r690027756>
> :
>
> > +        output_attentions (:obj:`bool`, `optional`):
> +            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
> +            tensors for more detail.
> +        output_hidden_states (:obj:`bool`, `optional`):
> +            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
> +            more detail.
> +        return_dict (:obj:`bool`, `optional`):
> +            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
> +"""
> +
> +
> ***@***.***_start_docstrings(
> +    "The bare gptj Model transformer outputting raw hidden-states without any specific head on top.",
> +    GPTJ_START_DOCSTRING,
> +)
> +class GPTJModel(GPTJPreTrainedModel):
>
> Again, I fear you misunderstand. You absolutely can perform inference on a
> 16 GB V100 GPU. At no point did I say that you needed 32 GB of memory to
> use this model.
>
> —
> You are receiving this because you commented.
> Reply to this email directly, view it on GitHub
> <https://github.com/huggingface/transformers/pull/13022#discussion_r690027756>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AA5MNWJ4WGL4LCRW3LVF7XTT5HRVPANCNFSM5BVGELLA>
> .
> Triage notifications on the go with GitHub Mobile for iOS
> <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675>
> or Android
> <https://play.google.com/store/apps/details?id=com.github.android&utm_campaign=notification-email>
> .
>


> Ultimately it's up to you - if it turns out that someone intending to utilize your PR is budget-constrained, they'll just be forced to end up making those changes locally by copy-pasting stuff from modeling_gpt2.py.

> Okay, well if that's the case, I'd love to see stats backing the claim that I misunderstand. Last I checked, fp32 weights for a 2.7B model itself is ~10 GB of disk space. So this 6B model you are creating a PR for will likely take more than 16 GB of disk space. So I do not see how you "absolutely can perform inference on a 16 GB V100 GPU" with this GPT-J 6B model. Please advise/correct my understanding since I'm still very much a newbie. Thanks!

You don't use fp32 weights for inference, you use BF16 weights without optimizer states for inference. BF16 weights without optimizer states for this model come out to 9GB. If you check out [the source repo](https://github.com/kingoflolz/mesh-transformer-jax), you can find a link to download them. As I've already said, you can find better optimized code at that repo, and if you're on a budget constraint you should not be using the `transformers` library.

As @EricHallahan says, the reality of large models is that you need large compute to run them. There is a fundamental limit to what you can do with a model that's too big to fit in your GPU, and that limit is best surpassed by buying a better GPU. We are well aware that this limits the accessibility of these models, but that's how to the world works. You can find other, smaller models that we have released [on our HF page](https://huggingface.co/EleutherAI) if our 6B model is outside your budgetary resources. This codebase also does not require that you use 6B parameters. You are welcome to use it to train a smaller model as well.

Hi @sgugger and @patil-suraj,

This PR's description says:

> the major design consideration was to make the configs compatible with GPT-2

I do not see the model classes being compatible though. For instance, `GPT2Model` supports `parallelize` (a very straightforward feature), `GPTJModel` does not at the moment. Authors of this PR don't seem keen on supporting this feature as you can see from the previous comments (including their edit history), but this is not their library, this is a community library. And usability by the community that makes this library famous and successful should indeed trump efficiency.

So should there be an addition of `parallelize` to this PR? What are your thoughts on this topic?

Thanks!

EDIT: I am fully aware I can fine-tune this model with standard gradient partitioning (ZeRO Stage 2) and at reasonably high TFLOPs/GPU with the usual old tricks, which I am not asking about. IMO the `parallelize` feature is immensely useful for already-trained models at deployment-time since it sets up a simple pipeline (unlike a formal `PipelineModule` sub-class expected for pipeline-parallel training) at inference-time for most consumers of this library, who may not be very technically savvy.

@g-karthik: I can back up the claims that it can be run in fp16 mode on any GPU with 16G of VRAM, but don't expect any large batching capabilities.

> Hi @sgugger and @patil-suraj,
> 
> This PR's description says:
> 
> > the major design consideration was to make the configs compatible with GPT-2
> 
> I do not see the model classes being compatible though. For instance, `GPT2Model` supports `parallelize` (a very straightforward feature), `GPTJModel` does not at the moment. Authors of this PR don't seem keen on supporting this feature as you can see from the previous comments (including their edit history), but this is not their library, this is a community library. And usability by the community that makes this library famous and successful should indeed trump efficiency.
> 
> So should there be an addition of `parallelize` to this PR? What are your thoughts on this topic?
> 
> Thanks!
> 
> EDIT: I am fully aware I can fine-tune this model with standard gradient partitioning (ZeRO Stage 2) and at reasonably high TFLOPs/GPU with the usual old tricks, which I am not asking about. IMO the `parallelize` feature is immensely useful for already-trained models at deployment-time since it sets up a simple pipeline (unlike a formal `PipelineModule` sub-class expected for pipeline-parallel training) at inference-time for most consumers of this library, who may not be very technically savvy.

This is a very uncharitable way to represent my assertion that I do not feel comfortable implementing a highly experimental feature that I didn't even know existed until you brought it up. I have no objection to it being implemented, and even explicitly invited you to do so.

If the HF team is interested in integrating this functionality across all of the `transformer` classes I have no objection to that whatsoever. However currently it has been implemented for 2 of the 66 (counting GPT-J) model classes.

@StellaAthena

My message says "Authors of this PR don't seem keen on supporting this feature as you can see from the previous comments (including their edit history)", where authors is **plural**. Eric is a co-author on this PR since I can see he has commits on this PR. I presume he is part of Eleuther-AI, since he was a co-author + you chose to apologize on his behalf and **then** invited me to push commits to your fork to add support for `parallelize`.

Your assertion was that you were not comfortable supporting `parallelize`, and your co-author Eric jumped the gun without bothering to look at the specific `parallelize` feature I referenced and made a broad assertion that parallelization "will not be considered" (which he later edited out after I responded, but I paste verbatim the original comment):

> GPT-J 6B easily fits in a 16 GB GPU for inference or 32 GB GPU for tuning (at FP16). CPU inference, while slow, absolutely remains an option and it isn't hard to get a server with 32 GB of memory. And besides, just because GPT-J 6B is the only model that exists today doesn't mean that other models cannot be created in the future at different scales. This is generic model class PR, not a GPT-J 6B PR.
Look, we understand that you are concerned with accessibility of large models. EleutherAI expects that those looking to use our large language models will need to find beefy hardware to run them. It is an unfortunate result of scaling, but it is something we cannot help with. You can try to do fancy optimizations and swapping weights in and out of memory but that is incompatible with the transformers design philosophy and out of scope for this PR.
If you want parallelization today, use Mesh Transformer JAX. Parallelization is out of scope for this PR, and will not be considered.

Given he was a co-author of your PR and the above observations, I think saying the "authors don't seem keen on supporting this feature" was actually most charitable. But to avoid clubbing you with your co-author, I shall rephrase to: "one author wasn't comfortable supporting this feature, and another author was just plain rude w.r.t this feature and gave out a flat no".

Anyway, I think I'm done w.r.t. this line of discussion, and shall wait until I hear from someone at Hugging Face (@sgugger or @patil-suraj) since I respect their design choices greatly and want to know their opinions on `parallelize` in GPT-J akin to GPT-2.

> @g-karthik: I can back up the claims that it can be run in fp16 mode on any GPU with 16G of VRAM, but don't expect any large batching capabilities.

@oborchers thanks for backing up their updated claim! The original (deleted, but reproduced by me since it's on email) claim was:

> Again, I fear you misunderstand. You absolutely can perform inference on a 16 GB V100 GPU. At no point did I say that you needed 32 GB of memory to use this model.

This does not state FP16 or BF16, which is why I presume the claim was quickly deleted.

One clearly cannot fit this GPT-J 6B model in a 16 GB GPU unless they use FP16 or BF16.

BUT, any consumer of a library as huge as Hugging Face transformers would prefer having the choice of whether they can/want to use FP32 or not. With the `parallelize()` and `deparallelize()` methods (supported in GPT-2), it would be possible for consumers to directly use FP32 weights on a 16 GB GPU because the model is split into pipeline stages for inference.

So, the argument here is simple. You have `modeling_gpt2.py` that `modeling_gptj.py` is supposedly meant to be riffing off of design-wise. The latter, however, currently does not support `parallelize()` and `deparallelize()`. If such support were added, power would lie in the hands of the consumer on whether or not to use FP32/FP16/BF16 for the use-case of their choice.

> Think the PR can be merged very soon :-)
> 
> We should probably try to focus on making the tests pass and then the only things that would be great to slightly adapt are:
> 
> * Remove the GPTJAttentionMixin and GPTJAttention class
> * Force the generation logits to be in fp32 so that the model can give good results in fp16 :-)
> 
> Thanks a lot for all the work on this already!

@EricHallahan is hoping to push the first change either today or tomorrow. Once he has, I'm expecting it'll be a couple minutes of work to fix the failing tests and ensure fp32 generation.

I see that `GPTJAttention().attn_dropout` and `GPTJAttention().masked_bias` are passed as parameters to `GPTJAttention()._attn()` rather than being referenced directly within `GPTJAttention()._attn()` like GPT-2. Similarly, `causal_mask` is calculated before `GPTJAttention()._attn()` and passed as a parameter while GPT-2 calculates it within `GPTJAttention()._attn()`. What does Hugging Face prefer?

The final failed test appears to be something technical about how the testing code was written and the removal of the `configs.rotary` argument. @patil-suraj, as you wrote most of the testing code could you take a look and see if you can spot how to fix it?

The traceback reads
```!bash
_______________ GPTJModelTest.test_gptj_model_past_large_inputs ________________
[gw0] linux -- Python 3.7.11 /usr/local/bin/python

self = <tests.test_modeling_gptj.GPTJModelTest testMethod=test_gptj_model_past_large_inputs>

    def test_gptj_model_past_large_inputs(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
>       self.model_tester.create_and_check_gptj_model_past_large_inputs(*config_and_inputs)

tests/test_modeling_gptj.py:386: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
tests/test_modeling_gptj.py:277: in create_and_check_gptj_model_past_large_inputs
    model = GPTJModel(config=config)
src/transformers/models/gptj/modeling_gptj.py:415: in __init__
    self.h = nn.ModuleList([GPTJBlock(config, layer_id=i) for i in range(config.n_layer)])
src/transformers/models/gptj/modeling_gptj.py:415: in <listcomp>
    self.h = nn.ModuleList([GPTJBlock(config, layer_id=i) for i in range(config.n_layer)])
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

self = GPTJBlock(
  (ln_1): LayerNorm((32,), eps=1e-05, elementwise_affine=True)
)
config = GPTJConfig {
  "activation_function": "gelu_new",
  "attention_probs_dropout_prob": 0.0,
  "attn_pdrop": 0.0,
  "bos_t...ts": true,
  "transformers_version": "4.10.0.dev0",
  "type_vocab_size": 16,
  "use_cache": true,
  "vocab_size": 99
}

layer_id = 0

    def __init__(self, config, layer_id):
        super().__init__()
        inner_dim = config.intermediate_size if config.intermediate_size is not None else 4 * config.n_embd
        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
>       self.attn = GPTJAttention(config, layer_id)
E       TypeError: __init__() takes 2 positional arguments but 3 were given

src/transformers/models/gptj/modeling_gptj.py:275: TypeError
```

Only two more tests to fix :-)

```
FAILED tests/test_modeling_gptj.py::GPTJModelTest::test_gptj_model_att_mask_past
FAILED tests/test_modeling_gptj.py::GPTJModelTest::test_gptj_model_past_large_inputs
```

@StellaAthena @EricHallahan those tests are pretty `transformers` specific and can be quite complex to fix - let me know if you want me to go into the PR to take a look :-)

> Only two more tests to fix :-)
> 
> ```
> FAILED tests/test_modeling_gptj.py::GPTJModelTest::test_gptj_model_att_mask_past
> FAILED tests/test_modeling_gptj.py::GPTJModelTest::test_gptj_model_past_large_inputs
> ```
> 
> @StellaAthena @EricHallahan those tests are pretty `transformers` specific and can be quite complex to fix - let me know if you want me to go into the PR to take a look :-)

Yeah that would be great! I would love it if you could take a look at what we are missing.

Eureka! [The result of the calculation intended for `attention_mask` was placed into `global_attention_mask` instead.](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gptj/modeling_gptj.py#L482-L501) The code passes both of the failing tests in question after replacing that block of code with [the corresponding block from GPT-2](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gpt2/modeling_gpt2.py#L696-L713). I'll integrate the change in a little bit after I do a little more testing.

@patrickvonplaten @sgugger @patil-suraj Looks like Eric saved the day! Let us know if there’s anything you’d like changed before it goes live.

Easy there @StellaAthena, we still need to verify that the slow tests pass. However, I am optimistic that we should be able to have this merged soon.

> Easy there @StellaAthena, we still need to verify that the slow tests pass. However, I am optimistic that we should be able to have this merged soon.

I’m under the impression that they need to manually approve the slow tests, no? That’s the “1 workflow awaiting approval” right?

I just ran them myself and they passed for me (after specifying `use_auth_token` in every call to `.from_pretrained()`), but I don't know if they will pass here.

I have a few matters that I think we should discuss/resolve before we consider merging:
1. I'll ask again because I haven't received a response yet: I see that [`GPTJAttention().attn_dropout` and `GPTJAttention().masked_bias` are passed as parameters to `GPTJAttention()._attn()`](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L233-L234) rather than being [accessed directly within the method like GPT-2](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt2/modeling_gpt2.py#L187-L194). Similarly, [`causal_mask` is calculated before `GPTJAttention()._attn()` and passed as a parameter](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L225-L226) while [GPT-2 calculates it within `GPT2Attention()._attn()`](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt2/modeling_gpt2.py#L185-L186). What does Hugging Face prefer? I think we should adapt the GPT-J implementation to be more like the GPT-2 implementation in this respect.
2. To resolve [a warning regarding `torch.where` deprecating uint8 condition tensors](https://github.com/pytorch/pytorch/blob/d7d399f3dfc780f3e49bcffe45694fb04e5db637/aten/src/ATen/native/TensorCompare.cpp#L330), [GPT-2 casts to bool after slicing `bias`](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt2/modeling_gpt2.py#L186). I resolved the same warning by [casting the contents of the entire `bias` buffer at initialization](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L79-L81). This seems to work fine, but if we need to change this for some reason I have not foreseen please tell me.
3. [GPT-2 calculates the scale factor for the attention weights and applies them within `GPT2Attention()._attn()`](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gpt2/modeling_gpt2.py#L180-L181) when [a bool config variable is set](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gpt2/modeling_gpt2.py#L147). [GPT-J used to have a dedicated buffer to store the scaling value](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gptj/modeling_gptj.py#L91), but I modified this [to remove the unneeded buffer](). This made the model loader stop complaining about buffers that were not initialized from the checkpoint. This seems to work fine (and if it does I think we can remove [the check that ensures that `scale_attn` is initialized](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L152)), but if we need to change this for some reason I have not foreseen please tell me.
4. I note that the weights currently staged on Model Hub are stored in [half precision (binary16)](https://en.wikipedia.org/wiki/Half-precision_floating-point_format), while the original released checkpoint was [bfloat16](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format). As this is an inherently lossy conversion unlike casting to [single precision (binary32)](https://en.wikipedia.org/wiki/Single-precision_floating-point_format), I feel it important to ask the level that we should adhere to the [`transformers` philosophy of <q>[providing] state-of-the-art models with performances as close as possible to the original models</q>](https://huggingface.co/transformers/philosophy.html). I have been assured by @StellaAthena and @kingoflolz that the difference in downstream performance between binary16 and bfloat16/binary32 is minimal with GPT-J 6B (and the evaluations they presented to me support this claim), but if we assume that this implementation will be used for academic research at some point in the future it seems odd to be manipulating the original checkpoint in a way that could modify downstream performance.  
The reason I bring this matter up here is firstly the fact that bfloat16 hardware is not as widespread/accessible as binary32 and binary16 hardware and secondly my assumption that a switch to a bfloat16 checkpoint would require changes in the implementation; If we decide that serving the checkpoint from Model Hub in bfloat16 is required to meet the goals of the `transformers` project, it is critical to ensure that it will be properly loaded/cast on platforms that do not support bfloat16 computation. I consider storing and serving the checkpoint in binary32 to be an unacceptable compromise to this conflict, as it would be double the size of the **11.7 GiB** binary16 checkpoint that is currently staged.
5. Would it be preferable to [include a tool that can convert the original checkpoint file to the Hugging Face format like GPT-Neo](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt_neo/convert_gpt_neo_mesh_tf_to_pytorch.py), or is that something that is out of scope of this PR?
5. It is unclear to me if [casting `query` and `key` to single precision (binary32) in `GPTJAttention()._attn()` is actually preventing an overflow when running the model in half precision (binary16)](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L145-L147). We should verify that this is required before merging.

Great work on this!

> 6\. It is unclear to me if [casting `query` and `key` to single precision (binary32) in `GPTJAttention()._attn()` is actually preventing an overflow when running the model in half precision (binary16)](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L145-L147). We should verify that this is required before merging.

The cast is originally there from the gpt-neo implementation. Did not find any difference, other than speed, when I evaluated not casting to fp32, so it doesn't seem to be necessary. However, when running the model on certain GPUs (e.g. P100) in bf16 mode, doing the matmul fails, while the model remains basically operable in bf16 mode if the cast is kept.

> Great work on this!
> 
> > 6. It is unclear to me if [casting `query` and `key` to single precision (binary32) in `GPTJAttention()._attn()` is actually preventing an overflow when running the model in half precision (binary16)](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L145-L147). We should verify that this is required before merging.
> 
> The cast is originally there from the gpt-neo implementation. Did not find any difference, other than speed, when I evaluated not casting to fp32, so it doesn't seem to be necessary. However, when running the model on certain GPUs (e.g. P100) in bf16 mode, doing the matmul fails, while the model remains basically operable in bf16 mode if the cast is kept.

@finetuneanon does a P100 have hardware support for bf16? Does other bf16 code run on it?

> I note that the weights currently staged on Model Hub are stored in half precision (binary16), while the original released checkpoint was bfloat16. As this is an inherently lossy conversion unlike casting to single precision (binary32), I feel it important to ask the level that we should adhere to the transformers philosophy of [providing] state-of-the-art models with performances as close as possible to the original models. I have been assured by @StellaAthena and @kingoflolz that the difference in downstream performance between binary16 and bfloat16/binary32 is minimal with GPT-J 6B (and the evaluations they presented to me support this claim), but if we assume that this implementation will be used for academic research at some point in the future it seems odd to be manipulating the original checkpoint in a way that could modify downstream performance.

Here’s the table of evaluation results @EricHallahan is talking about. All evaluations were done using [EleutherAI’s Eval Harness](www.github.com/eleutherai/lm-eval-harness):
| Task   | bf16     | fp32   |
| --- | --- | --- |
| Lambada | 0.683 | 0.697 |
| Winogrande | 0.648 | 0.653 |
| PiQA | 0.761 | 0.765 |
| HellaSwag | 0.661 | 0.661 |

> The reason I bring this matter up here is firstly the fact that bfloat16 hardware is not as widespread/accessible as binary32 and binary16 hardware and secondly my assumption that a switch to a bfloat16 checkpoint would require changes in the implementation; If we decide that serving the checkpoint from Model Hub in bfloat16 is required to meet the goals of the transformers project, it is critical to ensure that it will be properly loaded/cast on platforms that do not support bfloat16 computation. I consider storing and serving the checkpoint in binary32 to be an unacceptable compromise to this conflict, as it would be double the size of the 11.7 GiB binary16 checkpoint that is currently staged.

This is a HF philosophy question more than anything else, but for what it’s worth I would opt for both. My understanding is that the `transformers` library currently hides from the user the level of precision it is working with. I think that changing this so that user is able to control the precision being used is a very good idea. Due to HF’s no inheritance policy, it’s unclear to me how much work this would be to implement. I fear that every transformer needs to implement it individually though. We could sidestep this by implementing a precision flag that allows for fp16, bf16, and fp32 but doing so would likely incur significant technical debt.

> @finetuneanon does a P100 have hardware support for bf16? Does other bf16 code run on it?

It doesn't have hardware support, but it seems to get emulated in software for all operations other than the attn matmul, which causes a cublas error, from what I remember.

To allow picking the desired precision level as well as whether the model should get initialized directly on the GPU and whether to perform the matmul in fp32, I have added a few configuration parameters to my version somewhat recently:

* https://github.com/finetuneanon/transformers/blob/gpt-neo-localattention3-rp-b/src/transformers/models/gpt_neo/configuration_gpt_neo.py#L153-L155
* https://github.com/finetuneanon/transformers/blob/gpt-neo-localattention3-rp-b/src/transformers/models/gpt_neo/modeling_gpt_neo.py#L138-L161
* https://github.com/finetuneanon/transformers/blob/gpt-neo-localattention3-rp-b/src/transformers/models/gpt_neo/modeling_gpt_neo.py#L212-L215
* Also: L601, L810, L812

This is probably not clean enough to be included here.

> This is a HF philosophy question more than anything else, but for what it’s worth I would opt for both. My understanding is that the `transformers` library currently hides from the user the level of precision it is working with. I think that changing this so that user is able to control the precision being used is a very good idea. Due to HF’s no inheritance policy, it’s unclear to me how much work this would be to implement. I fear that every transformer needs to implement it individually though. We could sidestep this by implementing a precision flag that allows for fp16, bf16, and fp32 but doing so would likely incur significant technical debt.

The better solution would be to detect what the hardware supports and adapt accordingly; I think that we should serve the weights in bfloat16 and then cast them at load to the appropriate type (as long as this does not result in a significant usability penalty on resource-constrained machines).

To elaborate, we should prefer to load the model at its native bfloat16 unless the hardware fails to support it. After that decision resolving which format to cast to is tricky and we have a few options:

- We can simply call the problem out of scope let the end user cast to their preferred precision from bfloat16. This has issues with the "it just works" `transformers` mentality and would lead to problems when naively loading the model with a pipeline.
- We can cast the model to binary32 from bfloat16 and let the end user cast to their preferred precision (either binary16 or, in the case they need to override the detection, bfloat16). This has the issue of casting the weights to a larger type than they started as and hence potential OOM issues on resource constrained systems. We don't want to cast to binary32 without knowing the model will fit in memory, and ideally the heuristic should not cast to binary32 unless it is known to be the final type.
- We can detect if the hardware supports binary16 and then fallback on binary32 if that fails. This has the issue of the lossy conversion of the weights by default (which means that the end user may never know that they are working with a modified model that could be unsuitable for certain kinds of academic research). At first glance it may seem like the fallback to binary32 is another step that degrades the model, but because we are not loading the model until we decide on a type it isn't a problem here. This is the best solution for production use cases.

The issue with implementing something like these is that we don't know the device that the end user will be placing these models on, as models are placed on the device outside of the model constructor and `.from_pretrained()`. We can never know what features/formats are supported as they are out of scope of these methods. It would be possible to pass the intended device to the constructor but that would break with the "one consistent API" requirement. This seems to be a pretty tricky problem that is going to get more important to resolve as models continue to scale. @finetuneanon seems to have solved this, but not in a way that is acceptable to being integrated into `transformers`.

To assist in testing the code, I have made the EleutherAI GPT-J-6B checkpoint public. You can find it at https://huggingface.co/EleutherAI/gpt-j-6B, and the code should load from it by default.

> I think that we should serve the weights in bfloat16 and then cast them at load to the appropriate type (as long as this does not result in a significant usability penalty on resource-constrained machines).

As long as pytorch 1.9.1+ and CUDA 11.1+ are used and the attn matmul is cast to fp32, there should be no usability issues, but the model may be slower on GPUs that do not natively support bf16.

Also, I do not have any evaluation handy, but I think I observed at some point that casting the weights to fp16 and back to bf16 had very similar results to running in bf16 directly, which might indicate that the any runtime difference is mainly due to precision of activations and not so much of the weights. If this could be confirmed, casting the model to fp16 by default might be acceptable as well, as it could be cast back? Overall, casting the weights to fp16 does not cause any overflows and about 7k underflows where parameters became 0 where they weren't before.

@StellaAthena Thanks a lot to you and everyone involved here for all the work you've put into this.

I think there's a weird bug that I have verified doesn't happen with say distilgpt2 and may have to do with this model/PR specifically.

Basically, if you save the model and just load it again and start generating it will output gibberish, nothing close to what the model would output if you were to just load "EleutherAI/gpt-j-6B". Again, doesn't happen with distilgpt2, same exact setting.


```
import transformers
from transformers import AutoTokenizer, AutoConfig
from transformers import AutoModelForCausalLM
from transformers import Trainer


model_checkpoint = "EleutherAI/gpt-j-6B"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, use_fast=True)
config = AutoConfig.from_pretrained(model_checkpoint)
model = AutoModelForCausalLM.from_pretrained(model_checkpoint, config=config)

trainer = Trainer(
    model=model
)

trainer.save_model("gpt-j-test")
```

Now if you load the model that you saved and try to generate you'll see gibberish:

```
model = AutoModelForCausalLM.from_pretrained("gpt-j-test", config=config)
input_str = "Ciara is a software engineer"

inputs = tokenizer.encode(input_str, return_tensors="pt")
generated_text = model.generate(inputs.cuda(), max_length=133, repetition_penalty=1)


print(tokenizer.decode(generated_text[0]))
```

This will output something like:


```
Ciara is a software engineer product masc plates Strategymus Cal nervesaha Goddess gift loGW Bear RobPacific Cros DAMBalt surgical biotech permissions boxes fuckin customized vary Near Problems boxes Near lob automated overturn Cadillac giveauldron givethread viral valley giveauldron give 156 KILLprodu giveBUTaid Credits customized Strategy Ethernet week vary lo disregocusWA coasts viral Ni Processorlaughchart ol concentration Bend Large rescue overturn ol cornerback Bend OccupyStonesrcStonesrcStone give Starsonna customizedStonesrc Rosenstein CBD ashesStoneoddyStoneordered ol fond particip Bert Marion IN Dys Spe refugeeLikeLAN turbines disregLOC mouldLikeothal TiffLike Methodsiorsesi streak Eclipse Fuck space olLike Methods adjustmentsonna disreg 156
```


The normal output when you load "EleutherAI/gpt-j-6B" is something like:

```
Ciara is a software engineer at Google. She is also a member of the Google Cloud Platform team, and is responsible for the development of the Google Cloud Platform’s Kubernetes service.

In this post, Ciara will share her experience of working on Kubernetes at Google, and how she got started with Kubernetes.

Ciara’s journey with Kubernetes

I started working with Kubernetes in 2015, and I’ve been working on it ever since. I’ve been working on Kubernetes at Google for the last two years.
```





@emaconta Thanks for the report! We'll see of we can replicate this.

With saving:
> Ciara is a software engineerarlAndre library harvestedHEADrazelement 1899 camoufl regener photography reveal Mes- leukemia optional optional optional optional optional optional optionalusal Rex scissors sinkunker optional travers optional many optional stained optional Gi floor optionalrose arrangement optional Sugar WittNetMessage optionalJake optionalarov obtained optionalTorMissing optionalarov obtained optional bluesimmigrant obtained Slinventory Lantern optionalJake must columnist cannonput arrangementmbuds accurately cannon foot SlMonster cannon experiments testing arrangementarov obtaineddes Think BreweraezWalker obtainedarov obtained punk Ref McConnell HT symbolicavezbig liberals selections proves Witt Spread★★ ApplicantarovLocateddes selectionsMet ChickJesusMet focusdes selections recruits excessivelyMet objectmor vacanciesdes pubsheat enactCleanv Child

Without saving
> Ciara is a software engineer at Google. She is also a member of the Google Cloud Platform team, and is responsible for the development of the Google Cloud Platform’s Kubernetes service.
> 
> In this post, Ciara will share her experience of working on Kubernetes at Google, and how she got started with Kubernetes.
> 
> Ciara’s journey with Kubernetes
> 
> I started working with Kubernetes in 2015, and I’ve been working on it ever since. I’ve been working on Kubernetes at Google for the last two years.

Seems like a real issue. I'll look into it.

@emaconta It looks like `lm_head.weight` was blacklisted from being exported in checkpoints. Removing it from the blacklist looks to have solved the issue.

I meanwhile found out that attempting to load the model at binary16 (passing `torch_dtype=torch.half` to `.from_pretrained()`) crashes on CPU because [the construction of `GPTJAttention().scale_attn` utilizes `torch.sqrt()`](https://github.com/huggingface/transformers/blob/223bda109af8ceb36530fd6fae65daca48161516/src/transformers/models/gptj/modeling_gptj.py#L95). We'll have to fix that somehow.

I note that [`GPTJMLP()` still uses the names `c_fc` and `c_proj` (names that imply it uses convolutional layers when it uses `torch.nn.Linear`)](https://github.com/huggingface/transformers/blob/ad567a9e260eeea0c5e9878f0a2752f22719df19/src/transformers/models/gptj/modeling_gptj.py#L255-L256). We should change this for the purpose of readability.

@amantalion Have you tried loading the model at binary16/float16/half-precision? The Tensor Cores on the T4 do not include hardware for bfloat16. (bfloat16 computations are emulated on those cards.)

> Also, what's the speed that you are getting on your GPU that supports `bfloat16`?

A single A100 is finishing your benchmark in around 4.5 to 5 seconds, though that is a really poor point of comparison as an A100 has a TDP over three times larger than a T4 (250W >> 75W).

> It looks like it takes ~10.7 seconds to perform inference with GPU (T4 with 16GB VRAM) on the example that @emaconta provided. In comparison, it only takes 3.72 seconds with TPU (took from the sample Colab: https://colab.research.google.com/github/kingoflolz/mesh-transformer-jax/blob/master/colab_demo.ipynb).

I wouldn't look into that comparison too deeply without understanding the numerous differences between the conditions in which each benchmark was held. The codebase (this port vs. Mesh Transformer JAX), the framework (PyTorch vs. JAX), the platform (T4 vs. TPU v2-8), and the architecture (unparallelized on a single device vs. model-parallel across eight devices) are all different. TPU v2-8s are quite fast, especially with a heavily TPU-optimized codebase like Mesh Transformer JAX. If you are expecting the performance that the Colab notebook provides out of a T4, you unfortunately need to adjust your expectations.

> I have a few matters that I think we should discuss/resolve before we consider merging:
> 
> 1. I'll ask again because I haven't received a response yet: I see that [`GPTJAttention().attn_dropout` and `GPTJAttention().masked_bias` are passed as parameters to `GPTJAttention()._attn()`](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L233-L234) rather than being [accessed directly within the method like GPT-2](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt2/modeling_gpt2.py#L187-L194). Similarly, [`causal_mask` is calculated before `GPTJAttention()._attn()` and passed as a parameter](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L225-L226) while [GPT-2 calculates it within `GPT2Attention()._attn()`](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt2/modeling_gpt2.py#L185-L186). What does Hugging Face prefer? I think we should adapt the GPT-J implementation to be more like the GPT-2 implementation in this respect.
> 2. To resolve [a warning regarding `torch.where` deprecating uint8 condition tensors](https://github.com/pytorch/pytorch/blob/d7d399f3dfc780f3e49bcffe45694fb04e5db637/aten/src/ATen/native/TensorCompare.cpp#L330), [GPT-2 casts to bool after slicing `bias`](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt2/modeling_gpt2.py#L186). I resolved the same warning by [casting the contents of the entire `bias` buffer at initialization](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L79-L81). This seems to work fine, but if we need to change this for some reason I have not foreseen please tell me.
> 3. [GPT-2 calculates the scale factor for the attention weights and applies them within `GPT2Attention()._attn()`](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gpt2/modeling_gpt2.py#L180-L181) when [a bool config variable is set](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gpt2/modeling_gpt2.py#L147). [GPT-J used to have a dedicated buffer to store the scaling value](https://github.com/huggingface/transformers/blob/b6021cf0d9acc36fb96aaf3c7b457160f7f0b9d5/src/transformers/models/gptj/modeling_gptj.py#L91), but I modified this to remove the unneeded buffer. This made the model loader stop complaining about buffers that were not initialized from the checkpoint. This seems to work fine (and if it does I think we can remove [the check that ensures that `scale_attn` is initialized](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L152)), but if we need to change this for some reason I have not foreseen please tell me.
> 4. I note that the weights currently staged on Model Hub are stored in [half precision (binary16)](https://en.wikipedia.org/wiki/Half-precision_floating-point_format), while the original released checkpoint was [bfloat16](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format). As this is an inherently lossy conversion unlike casting to [single precision (binary32)](https://en.wikipedia.org/wiki/Single-precision_floating-point_format), I feel it important to ask the level that we should adhere to the [`transformers` philosophy of [providing] state-of-the-art models with performances as close as possible to the original models](https://huggingface.co/transformers/philosophy.html). I have been assured by @StellaAthena and @kingoflolz that the difference in downstream performance between binary16 and bfloat16/binary32 is minimal with GPT-J 6B (and the evaluations they presented to me support this claim), but if we assume that this implementation will be used for academic research at some point in the future it seems odd to be manipulating the original checkpoint in a way that could modify downstream performance.
>    The reason I bring this matter up here is firstly the fact that bfloat16 hardware is not as widespread/accessible as binary32 and binary16 hardware and secondly my assumption that a switch to a bfloat16 checkpoint would require changes in the implementation; If we decide that serving the checkpoint from Model Hub in bfloat16 is required to meet the goals of the `transformers` project, it is critical to ensure that it will be properly loaded/cast on platforms that do not support bfloat16 computation. I consider storing and serving the checkpoint in binary32 to be an unacceptable compromise to this conflict, as it would be double the size of the **11.7 GiB** binary16 checkpoint that is currently staged.
> 5. Would it be preferable to [include a tool that can convert the original checkpoint file to the Hugging Face format like GPT-Neo](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gpt_neo/convert_gpt_neo_mesh_tf_to_pytorch.py), or is that something that is out of scope of this PR?
> 6. It is unclear to me if [casting `query` and `key` to single precision (binary32) in `GPTJAttention()._attn()` is actually preventing an overflow when running the model in half precision (binary16)](https://github.com/huggingface/transformers/blob/d2c85a22abb7b09185363a67a270276465779182/src/transformers/models/gptj/modeling_gptj.py#L145-L147). We should verify that this is required before merging.

Hey @EricHallahan,

Thanks a lot for the very in-detail summary of the remaining TODO's. Here is my opinion:

1) Agree! I'll change this directly in the PR now (hope that's fine for you)
2) That's totally fine! The reason we created the mask in `torch.int8` for GPT2 was that back then we wanted to keep supporting PyTorch 1.1.0 which didn't support the `torch.bool` type. Now our minimum required version for torch is 1.3.0 -> so creating it directly in `torch.bool` is preferred. So this looks good to me!
3) Here we can remove the `if ...` check since the attention is always scaled -> I'll also adapt this in the PR directly.
4) That's a very important one here! We would prefer to actually store the weights in fp32 on the hub as all of our model weights are stored in fp32 and this is what the user expects. As you said storing it in fp16 is not a good idea because of the lossy conversion and since we want to provide the models as similar as possible to the original implementation as possible. The other option is to store them in bfloat16, which would be closer the original implementation as possible, but the problem is that lots of hardware doesn't support bfloat16 yet in PyTorch and we would like the model to work out of the box, *e.g.* without the user having to know that `.to(torch.float32)` is required to run the model on CPU. Experienced users that are aware of bfloat16 can then cast the checkpoint to `bfloat16` without loosing precision etc...Also all of our "officially supported" models are stored in float32 on the hub so I'd prefer to not make an exception here (*e.g.* all the T5 weights are also stored in fp32 even though the official weights are bf16). Let me know what you think here!
5) Yes such a conversion file would be nice! IMO this can be done in a follow-up PR though (happy to help on this as well).
6) Also a very important point. After discussing with @LysandreJik a bit, I think we should keep hard-coded casts to fp32 in the modeling code if necessary so that the model can run in float16 at the moment. By necessary we mean that if a model requires some weights/operation to be cast to `float32` in order to prevent over-/underflowing issues in `float16` then we should keep it in float32. If the model doesn't require the `float32` cast to run in float16, then we can remove it. Another example where this is used is T5's special layer norm Module: https://github.com/huggingface/transformers/blob/f689743e7454b93f6cab4343026de03fa530bfb9/src/transformers/models/t5/modeling_t5.py#L242. We don't support bfloat16 officially in the library yet, so I would like to base the decision here first just on float32 vs. float16. Once `bfloat16` is officially supported, we will do a pass on all models and make sure they work correctly in bfloat16. Let me know what you think here!
7) RE: the naming of `c_proj` - I agree with you we should rename the variables here since there is no convolution layer anymore - I'll also adapt this in this PR 

I have applied the changes as mentioned in 1.), 3.), 4.) and 6.), 7.) and also changed the official weights accordingly here: https://huggingface.co/EleutherAI/gpt-j-6B/commit/a16214f67cb321168930e17ee9a508794721bc3d . Note that the weights are in fp32 now which is why the model weights file is twice as big. 

The output logits are also now "hard"-casted to `float32` to make sure generation works correctly in fp16 as stated by @patil-suraj.

Since the weights are stored now in fp32 on the hub, one needs to load the model with `torch_dype=torch.float16` if there is not enough RAM on the GPU (I also left some notes in the docs about this).

```python
from transformers import AutoTokenizer, GPTJForCausalLM
import torch
  
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", torch_dtype=torch.float16).to("cuda")

def gen(input_str):
    input_ids = tokenizer(input_str, return_tensors="pt").input_ids.to(model.device)
    gen_out = model.generate(input_ids, do_sample=True)
    print(tokenizer.batch_decode(gen_out))

gen("Today is a nice")
```

Let me know if this is ok for you @EricHallahan @StellaAthena or if you would like me to revert a commit or whether we should do anything else before merging!

IMO, the PR is good to merge (cc @LysandreJik) 

I fear there’s a little confusion about the model weights I want to clarify.

The original model training occurred in fp32. Because fp32 is so memory intensive, we created bf16 weights with the expectation that they would be used for inference. Since the transformers package puts a strong emphasis on accessibility, we felt that the fp32 and bf16 weights would not be appropriate for the HF model. The fp32 weights are too large to use on a consumer GPU, and many GPUs don’t support bf16 yet. Since bf16 isn’t commonly supported, I made the decision to upload the weights cast as fp16 to follow HF philosophy of valuing accessibility. When @patrickvonplaten moved the weights to fp32, he didn’t use the real fp32 weights: he took the fp16 weights and cast them as fp32. These weights, despite being listed as fp32, as no better than using fp16 and is really just padded with 11 GB of null.

As a result, there are four sets of possible weights:
1. True fp32 weights. These are inaccessible to most people due to their size, and will not even download on most consumer systems
2. bf16 weights: these are the current widespread best practices for working with a model like GPT-J, but most GPUs don’t support bf16 yet.
3. fp16 weights: these are widely accessible and have nearly identical performance to bf16 for inference in our testing.
4. The currently uploaded weights: these have the precision of fp16 while taking up the space of fp32. These should not be used.

I think it makes the most sense to have two models: `eleutherai/gpt-j-6B` and `eleutherai/gpt-j-6B-fp32` or something like that. People who have systems that have enough space for the fp32 weights can use them, but the overwhelming majority of people will instead use the smaller fp16 weights. fp16 weights can be losslessly cast to bf16 by people who have hardware support for it.

> I fear there’s a little confusion about the model weights I want to clarify.
> 
> The original model training occurred in fp32. Because fp32 is so memory intensive, we created bf16 weights with the expectation that they would be used for inference. Since the transformers package puts a strong emphasis on accessibility, we felt that the fp32 and bf16 weights would not be appropriate for the HF model. The fp32 weights are too large to use on a consumer GPU, and many GPUs don’t support bf16 yet. Since bf16 isn’t commonly supported, I made the decision to upload the weights cast as fp16 to follow HF philosophy of valuing accessibility. When @patrickvonplaten moved the weights to fp32, he didn’t use the real fp32 weights: he took the fp16 weights and cast them as fp32. These weights, despite being listed as fp32, as no better than using fp16 and is really just padded with 11 GB of null.
> 
> As a result, there are four sets of possible weights:
> 
> 1. True fp32 weights. These are inaccessible to most people due to their size, and will not even download on most consumer systems
> 2. bf16 weights: these are the current widespread best practices for working with a model like GPT-J, but most GPUs don’t support bf16 yet.
> 3. fp16 weights: these are widely accessible and have nearly identical performance to bf16 for inference in our testing.
> 4. The currently uploaded weights: these have the precision of fp16 while taking up the space of fp32. These should not be used.
> 
> I think it makes the most sense to have two models: `eleutherai/gpt-j-6B` and `eleutherai/gpt-j-6B-fp32` or something like that. People who have systems that have enough space for the fp32 weights can use them, but the overwhelming majority of people will instead use the smaller fp16 weights. fp16 weights can be losslessly cast to bf16 by people who have hardware support for it.

Thanks for your comment here @StellaAthena! 

The current weights are indeed not ideal, as they are just padded fp16 weights.

Would it be ok for you however to upload the "true" fp32 weights? People can quite easily convert the weights from fp32 to fp16 under the hood via:

```python
model = GPTJForCausalLM.from_pretrained('...', torch_dtype=torch.float16)
```
and then move the casted weights to GPU. From our side this would be the preferred behavior - would that be ok for you?

@patrickvonplaten Can we do both fp32 and fp16 models? I care a lot about Google Colab because it's the way that most people have access to the necessary GPUs to use these models and the fp32 weights won't download into Colab's RAM.

I am not only concerned about Colab either: I expect that a common system configuration for deploying GPT-J 6B at home with `transformers` will be an RTX 3090 with 16 GiB of system memory. Attempting to load GPT-J 6B at single precision will OOM before it can be cast and loaded onto the GPU, despite it being possible to run GPT-J 6B on this system configuration otherwise.

We actually picked the size of GPT-J specifically to allow it to be loaded onto Google Colab or a RTX 3090 with enough room left over to fine-tune it at what we judged to be a reasonable rate. ~11 GiB is the largest you can make a model that will be functional on consumer GPUs under current technology. I really think that requiring people to download it in fp32 will be disastrous for adoption.

I share the above concerns about accessibility.

As an example use case, I use a finetuned GPT-J in my social media bot.  For users who are curious about how the bot works, I have a [Colab inference demo](https://colab.research.google.com/drive/1pdRISGaypV8DwolvD3uv77vlmMX0xZNL?usp=sharing).  The demo currently uses @finetuneanon’s fork and an fp16 checkpoint, the same checkpoint I use for inference in the bot itself.  It works fine on an ordinary Colab instance with a T4 GPU and 12 GB system RAM.

Most of my bot’s curious users don’t have access to more specialized hardware.  It’s great to be able to provide these people with direct access to such a powerful model.

I see! It's indeed very nice to be able to run the model in a Google Colab for which we need half-precision. I think we should opt then for two models. 

What do you think about adding a `gpt-j-6b-fp16` checkpoint in float16 precision in addition to the `gpt-j-6b` checkpoint that will be kept in full fp32 precision?

> I see! It's indeed very nice to be able to run the model in a Google Colab for which we need half-precision. I think we should opt then for two models.
> 
> What do you think about adding a `gpt-j-6b-fp16` checkpoint in float16 precision in addition to the `gpt-j-6b` checkpoint that will be kept in full fp32 precision?

I liked the idea of making the fp32 on model hub the real fp32 and not the bf16 padded to fp32, If EleutherAI is willing to do that.

> I see! It's indeed very nice to be able to run the model in a Google Colab for which we need half-precision. I think we should opt then for two models.
> 
> What do you think about adding a `gpt-j-6b-fp16` checkpoint in float16 precision in addition to the `gpt-j-6b` checkpoint that will be kept in full fp32 precision?

I would prefer to do `gpt-j-6B` and `gpt-j-6B-fp32` for two reasons:
1. It makes more sense to me that the version with the simplest name be the most used verison.
2. In practice this is a breaking change. Yes, it hasn’t been officially released yet, but this model has been in widespread use for a while. Since I made the checkpoint public it has been downloaded over 1500 times, and since the weights were changed to fp32 I have already fielded questions from five individuals and one company as to why their model suddenly stopped working.

More generally, I am surprised and confused by this proposal because it makes me feel like I don’t understand HF’s philosophy. But hey, it’s your package. I would be happy to do it either way – my main priority is having a fp16 version.

The philosophy behind having the "main" model in fp32 is that it's as close as possible to the original model. We don't have any models saved in fp16 on the model hub, but always in fp32 so that the user has the freedom to decide which dtype suits better. Also just to make sure I'm not on the wrong page here - as I understand the original weights: https://the-eye.eu/public/AI/GPT-J-6B/step_383500.tar.zstd are in fp32 format no? 

I understand that this is different here because the full model fp32 cannot be loaded, *e.g.* in a Google Colab though and that the model is particularly large -> therefore I think it's a good idea to also provide the lossy fp16 weights, but I don't think it justifies having the default weights being a compressed version of the true original weights.

In the end, it's totally up to you what is better as the model is under your organization. So we are happy to go along with having fp16 as the default if you prefer - it would however be great to also have the fp32 version exposed so that there is access to the full precision. My goal was mainly to speed-up the merging process here :-) 

On a slightly different aspect, it might also be nicer to actually just have one "gpt-j-6b" model instead of two. We could use a git branch to differentiate between `fp16` and `fp32`. So with the git-lfs version control we could either save the fp16 or fp32 weights under a branch called `fp16` or `fp32` and then have the main branch be the other version. 

Right now fp32 is the main version, so doing the following:

```python
from transformers import GPTJForCausalLM

model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")
```

will load the 23GB file 

while 

```python
from transformers import GPTJForCausalLM

model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="fp16")
```

will load the fp16 12GB file that can be used in a Google colab because I added a branch `fp16` to the model here: https://huggingface.co/EleutherAI/gpt-j-6B/tree/fp16 . I think this is a bit nicer instead of having two different models on the hub, no? 

=> So to summarize I think we have 4 options here:

1. `fp16` is the main model and we create a branch `fp32`:

```python
from transformers import GPTJForCausalLM

model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")  # <- for fp16
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="fp32") # <- for original full weights
```

2. `fp32` is the main model and we create a branch `fp16` (just leave the branch we have now)

```python
from transformers import GPTJForCausalLM

model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B")  # <- for original full weights
model = GPTJForCausalLM.from_pretrained("EleutherAI/gpt-j-6B", revision="fp16")  # <- for fp16
```

That's more or less what we have now

3. Create two models with `EleutherAI/gpt-j-6B` being in fp16 and `EleutherAI/gpt-j-6B-fp32` being in fp32
4. Create two models with `EleutherAI/gpt-j-6B` being in fp32 and `EleutherAI/gpt-j-6B-fp16` being in fp16

Maybe you can decide and then I can help to upload the weights correctly. 

Yes seconding what @patrickvonplaten said, @StellaAthena it's your model so it's your call! 

Let us know which option you prefer 

I just finished modifying the conversion script so that it can spit out `transformers` checkpoints at any precision from Mesh Transformer JAX checkpoints, so it should be a pretty painless process to upload newly processed weights when we decide what we want to do.

My opinion is that the only sensible way of maintaining multiple checkpoints of the same model is with multiple branches in a single model repo and exposing the choice to the end user via the revisions system.

I think my preferred configuration is to name the default branch *`float32`* for the single precision checkpoint and create a second branch called *`float16`* for the half precision checkpoint. (There is nothing stopping us from creating a third branch for *`bfloat16`*, but I expect demand for it to be very low and so is probably not worth the effort.) I think *`float32`* and *`float16`* are the obvious choice for the branch name pattern as [they mirror the `torch_dtype` of each model](https://huggingface.co/transformers/main_classes/configuration.html?highlight=torch_dtype). I also recommend committing freshly converted checkpoints to these branches to verify they are as clean as possible before we merge this PR.

We will need to make it abundantly clear that multiple branches exist within the model cards. I am thinking of placing a identifier directly under the first level heading so that it is visible above the fold on Model Hub, plus a short paragraph with an explanation of when to use that branch over others (drafts below). I also suggest more detailed code examples than normal that specifically demonstrate how to mix the model `revision` and `torch_dtype` for the purpose of clarifying their usage and interaction.

<figure>
<dl>
<dt><em><code>float32</code></em></dt>
<dd>This checkpoint of GPT-J 6B is stored in <a href="https://en.wikipedia.org/wiki/Single-precision_floating-point_format">single precision</a> and is most suitable for academic and research applications that require as close to original downstream performance. This 23.4 GiB checkpoint can be readily cast to lower-precision formats such as half precision and bfloat16 after loading. Given that there is no statistically significant difference in downstream performance when GPT-J 6B is run with reduced precision, it is recommended to use the alternative half precision checkpoint in prototyping and production applications.</dd>
<dt><em><code>float16</code></em></dt>
<dd>This checkpoint of GPT-J 6B is stored in <a href="https://en.wikipedia.org/wiki/Half-precision_floating-point_format">half precision<a> and is most suitable for prototyping and production applications where speed and resource-constraints are critical factors. This 11.7 GiB checkpoint can be readily cast to other floating-point formats for use on hardware that does not support half precision, a usage that saves both time and storage space over a higher-precision checkpoint. Half precision comes at the cost of a slightly different performance in downstream tasks, and it is recommended to use the alternative single precision checkpoint in the academic and research applications where this is not acceptable.</dd>
</dl>
</figure>

I really like everything @EricHallahan said. And I think @patrickvonplaten's idea is brilliant: if I had known that was an option I would have proposed it.

Great! @StellaAthena @EricHallahan let me know if you need any help uploading the checkpoints to the branches! 

@EricHallahan - could you send me a link to your updated conversion script? I'v tried using this one: https://github.com/kingoflolz/mesh-transformer-jax/pull/108 . The converted output is a dict with all layers saved as individual files. Is this the official conversion script? 

Just a note that for organizing models outside the hugging face cache, it is more convenient to have subfolders or separate repos for different content, because git-lfs can be very slow filtering many gigabytes when switching branches.  Not planning on arguing the point, just making sure the use-case is shared.

@patrickvonplaten The script currently within Mesh Transformer JAX builds a split checkpoint in the format used by the @finetuneanon fork of `transformers`, and I needed to heavily modify it to generate checkpoints in the HF format. @kingoflolz has already asked me to make a PR to update the script.

If we feel that the solution I outlined is the best solution, I can put that plan into action and update the repo on Model Hub. Maybe vote 👍/👎 on this?

@g-karthik I have ported over the experimental parallelization code from GPT-2 into GPT-J. I wouldn't personally recommend using it to anyone unless they need to, but it should work in a pinch when a better solution is unavailable.

(I note that this bears no resemblance to the implementation of model pararallelism in Mesh Transformer JAX and should not be thought as an equivalent implementation or replacement for that implementation.)

@EricHallahan: Thank you very much for the awesome work on the issue!! there is one thing to remark regarding the 16 VRAM configuration:

As far as I can tell, even with the floating point revision set correctly (and also when only the local files are fp16, local_files_only=True), the model will still be loaded in float32 until model.half() is called, thus requiring the 23G RAM to be available before the model.half() and before the model.to(device) is called.

By extension this means, that a text generation pipeline cannot be loaded with device=0 in a VRAM<23G setting, as .half() isn’t called automatically anywhere. In this case the model must be loaded, .halved(), and then passed to the pipeline via the argument.

Correct me if this observation is wrong. Is there any way of loading/moving the model iteratively to GPU so that the 23G RAM limitation can be circumvented, similar as done in @finetuneanon repository? (Probably out of scope for this very PR, but likely a problem for larger models in general in the future). Presumably this can be done using the state dict, but I‘m not deep enough into the inner working to judge this.

Also tagging @patrickvonplaten  

@oborchers Yes, we have had multiple people test this via Colab and they have reported the same issue.

I have verified that choosing the `"float16"` revision loads the model at float32. I don't understand why it doesn't load the model at half precision, especially because [I explicitly set `"torch_dtype": "float16"` in the model config on the *`float16`* branch](https://huggingface.co/EleutherAI/gpt-j-6B/commit/b380a4785c3184468823bb0bd98dcd4453c9d604). Maybe I am interpreting the fuctionality of that config parameter wrong, but my understanding is that it explicitly tells the model loader to use the specified type as the default.
(I also want to make sure to point out that naively loading the model is not particularly useful at this time, as the Model Hub repo still has an extraneous *`main`* branch that I have been unable to remove and replace with the *`float32`* branch.)

The multi-part loading scheme used by the @finetuneanon fork was purposefully built to bypass the suboptimal way that `transformers` loads checkpoints so that resource-constrained systems could load GPT-Neo (and later GPT-J) without running out of memory. In order to meet the requirements for integration into `transformers` we had to adapt that code to instead use the existing single-file checkpoint format. It is up to the `transformers` maintainers to consider an alternative/optimized checkpoint loading pipeline, and I assume that such a system would need a separate PR considering the changes probably needed to `PretrainedModel`.

> @oborchers Yes, we have had multiple people test this via Colab and they have reported the same issue.
> 
> I have verified that choosing the `"float16"` revision loads the model at float32. I don't understand why it doesn't load the model at half precision, especially because [I explicitly set `"torch_dtype": "float16"` in the model config on the _`float16`_ branch](https://huggingface.co/EleutherAI/gpt-j-6B/commit/b380a4785c3184468823bb0bd98dcd4453c9d604). Maybe I am interpreting the fuctionality of that config parameter wrong, but my understanding is that it explicitly tells the model loader to use the specified type as the default.
> (I also want to make sure to point out that naively loading the model is not particularly useful at this time, as the Model Hub repo still has an extraneous _`main`_ branch that I have been unable to remove and replace with the _`float32`_ branch.)
> 
> The multi-part loading scheme used by the @finetuneanon fork was purposefully built to bypass the suboptimal way that `transformers` loads checkpoints so that resource-constrained systems could load GPT-Neo (and later GPT-J) without running out of memory. In order to meet the requirements for integration into `transformers` we had to adapt that code to instead use the existing single-file checkpoint format. It is up to the `transformers` maintainers to consider an alternative/optimized checkpoint loading pipeline, and I assume that such a system would need a separate PR considering the changes probably needed to `PretrainedModel`.

Thanks a lot for the detailed message here. What we currently do in `.from_pretrained(...)` is definitely suboptimal if one is sure that all the loaded parameters are correct and complete. What happens under-the-hood is that:

1. A random model with the correct configuration is instantiated meaning all layers are randomly initialized as defined in the config. Random initialization is always happening in fp32.
2. Then the `state_dict` is loaded (the correct weights)
3. All layers of the random model are compared to the "real" layers that can be found in the state dict
4. All weights of the layers that are present in the state dict are dropped and overwritten by the state dict
5. All layers are casted to the correct dtype (as defined by `torch_dtype`)

=> we have this logic mainly for models like BERT for which one would load the "base"-model and than add a randomely initialized head for the specific downstream task. It becomes quite clear however that this make less sense for GPT-like models. 
There is an open issue to solve this problem: https://github.com/huggingface/transformers/issues/12274 but I don't think it'll be that easy to solve. 
If ok with you guys (@EricHallahan @StellaAthena) we would merge the PR and then try to fast-track the issue - what do you think?




FYI: Model acceleration for GPT-J via deepspeed in the making: https://github.com/microsoft/DeepSpeed/issues/1332

## Important

We will merge GPT-J now to master. Note that at the moment GPT-J **cannot** be run on a free google colab GPU since loading the model weights in fp16 requires too much CPU RAM. At the moment one needs at least **26 GB** of CPU RAM in order to load GPT-J in fp16-precision. We are working on fixing the problem so that in a next step one can load GPT-J with just 12 GB of CPU of RAM.

I feel the need to reiterate that there [remains a redundant *`main`* branch in Model Hub](https://huggingface.co/EleutherAI/gpt-j-6B/tree/main) that is neither [the true single precision checkpoint found in *`float32`*](https://huggingface.co/EleutherAI/gpt-j-6B/tree/float32) nor [the half precision checkpoint found in *`float16`*](https://huggingface.co/EleutherAI/gpt-j-6B/tree/float16). This means that naive usage (i.e. not specifying `revision="float32"` or `revision="float16"`) will not download the proper checkpoints.

> reiterate

Thanks for letting me know - is it ok if I put the "correct fp32 weigts" in the main branch for now? Or do you prefer "fp16"? Both are fine with us :-) Think we can't completely delete the "main" branch for now (cc @LysandreJik)

> Think we can't completely delete the "main" branch for now

That is my understanding.

>  is it ok if I put the "correct fp32 weigts" in the main branch for now? Or do you prefer "fp16"? Both are fine with us :-)

Putting the single precision weights in *`main`* should be fine for now.

> > Think we can't completely delete the "main" branch for now
> 
> That is my understanding.
> 
> > is it ok if I put the "correct fp32 weigts" in the main branch for now? Or do you prefer "fp16"? Both are fine with us :-)
> 
> Putting the single precision weights in _`main`_ should be fine for now.

+1 this

Ok great - just uploaded the correct weigths to "main". You can see that the sha256 between "main": https://huggingface.co/EleutherAI/gpt-j-6B/blob/main/pytorch_model.bin and "float32" https://huggingface.co/EleutherAI/gpt-j-6B/blob/float32/pytorch_model.bin match now :-)