Megatron GPT2 not compatible with transformers
Author: hwijeen
Pull request: no

**TLDR:**
Current implementation detail of attention in GPT2 is different from Megatron-LM GPT2, raising compatibility issue. Since the required change is quite minimal, how about changing transformers's implementation to follow Megatron-LM's?


**Quite detailed exploration:**
I used the [script](https://github.com/huggingface/transformers/blob/master/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py) to convert Megatron-LM GPT2 into transformers GPT2.
The script itself works well, but the generation results with the converted checkpoint show that there is something wrong.
```python
from transformers import GPT2LMHeadModel
from transformers import GPT2Tokenizer

model = GPT2LMHeadModel.from_pretrained(".")
tok = GPT2Tokenizer.from_pretrained("gpt2")

input_ids = tok.encode("How are you doing these days?" , return_tensors="pt")
gen_tokens = model.generate(
    input_ids,
    max_length=150,
    min_length=30,
    do_samples=True,
    num_beams=1,
    temperature=0.9,
    top_p=0.8,
    top_k=0,
    repetition_penalty=5.0,
)
print(tok.batch_decode(gen_tokens)[0])
# How are you doing these days?` `'' ''",'',",..., and '´','',", – ---''truvesabendabyaruobirockeysludboatmanlongsoceyeand how nobody anybody anything no way or what why not with all of theofoldways is by showing that if it's going to skip out onansonsarry hold skator who wonder when they're shown their a one two those well saying so there but behind them only ones don doubt wonders just like he his soon Johnny Louous ever after asking holding good say hey even though its still yet we know as long for now then may be in show off telling said terking something which will start next thing already says
```

I looked into the problem and found out that the logits produced by Megatron-LM GPT-2 and transformers GPT-2 are different.
Careful debugging revealed that it stems from the different way of implementing the attention mechanism.

```python
# Megatron-LM
# splits head first
    225         # [sq, b, (np * 3 * hn)] --> [sq, b, np, 3 * hn]
    226         new_tensor_shape = mixed_x_layer.size()[:-1] + \
    227             (self.num_attention_heads_per_partition,
    228              3 * self.hidden_size_per_attention_head)
    229         mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)

# and then splits key,query,value
    230
    231         # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
    232         (query_layer,
    233          key_layer,
    234          value_layer) = mpu.split_tensor_along_last_dim(mixed_x_layer, 3)

# result
query_layer.norm() # tensor(110.7500, device='cuda:0', dtype=torch.float16, grad_fn=<CopyBackwards>)
key_layer.norm() # tensor(125.3125, device='cuda:0', dtype=torch.float16, grad_fn=<CopyBackwards>)
value_layer.norm() # tensor(42.9688, device='cuda:0', dtype=torch.float16, grad_fn=<CopyBackwards>)
```

```python
# transformers
# splits key, query, value first
     242             query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)

# and then splits head

    244         query = self._split_heads(query, self.num_heads, self.head_dim)                                                                                                      
    245         key = self._split_heads(key, self.num_heads, self.head_dim)                                                                                                          
    246         value = self._split_heads(value, self.num_heads, self.head_dim)

# result
query.norm() # tensor(103.4170, device='cuda:0', grad_fn=<CopyBackwards>)
key.norm() # tensor(102.9513, device='cuda:0', grad_fn=<CopyBackwards>)
value.norm() # tensor(92.2508, device='cuda:0', grad_fn=<CopyBackwards>)
```

When I modified transformers GPT2 implementation to match Megatron-LM, the generation looks correct.
```python
# before
How are you doing these days?` `'' ''",'',",..., and '´','',", – ---''truvesabendabyaruobirockeysludboatmanlongsoceyeand ....

# after
How are you doing these days?, even further all," not a as to also 3 both every if [ always and ever ....
```

The generation results does not look super convincing, probably because it is not a very capable model (only ~300M parameters). When I experimented with my own LM checkpoint of 1.3B, the generation is only sensible when I made the modification.
Plus, I did a sanity check and confirmed that the logits from the Megatron-LM GPT-2 and transformers GPT-2 are the same only when I made the modification [here](https://github.com/huggingface/transformers/blob/61c506349134db0a0a2fd6fb2eff8e29a2f84e79/src/transformers/models/gpt2/modeling_gpt2.py#L242).
```python
# modify to follow Megatron-LM's implementation
else:
   # query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
   attn_out = self.c_attn(hidden_states)
   attn_out = self._split_heads(attn_out, self.num_heads, self.head_dim * 3)
# query = self._split_heads(query, self.num_heads, self.head_dim)
# key = self._split_heads(key, self.num_heads, self.head_dim)
# value = self._split_heads(value, self.num_heads, self.head_dim)
query, key, value = attn_out.split(self.head_dim, 3)
```

How about merging this change into master branch? I would be happy to make a PR.

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.6.1
- Platform: Linux-4.18.0-25-generic-x86_64-with-glibc2.10
- Python version: 3.8.5
- PyTorch version (GPU?): 1.8.1+cu102 (True)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help
@LysandreJik @jdemouth 
<!-- Your issue 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: @Rocketknight1

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

Model hub:

- for issues with a model report at https://discuss.huggingface.co/ and tag the model's creator.

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

 -->

## Information

Model I am using (Bert, XLNet ...): GPT2 (converted from Megatron-LM)

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [ ] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [ ] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. Load [gpt-2 checkpoint](https://ngc.nvidia.com/catalog/models/nvidia:megatron_lm_345m) with Megatron-LM and do a forward pass.
2. Convert the checkpoint into transformers-compatible format using the [script](https://github.com/huggingface/transformers/blob/master/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py), load the model with transformers and do the same forward pass.
3. The logits are not the same.

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
<!-- A clear and concise description of what you would expect to happen. -->
1. Generation results of converted GPT2 are sensible.
2. Logits from Megatron-LM GPT2 and transformers GPT2 are the same.

Comments:
This is the code to check if the two gpt2s produce the same logits. Hope this helps and I would be happy to provide further information if requested.
* megatron-LM
```python
import torch
from megatron import get_args, get_tokenizer, initialize_megatron, mpu
from megatron.model import GPTModel
from megatron.training import get_model
from megatron.checkpointing import load_checkpoint
from megatron.utils import get_ltor_masks_and_position_ids

def initialize():
    model_path = "/workspace/my_gpt3_1.3B_mp1"
    sys.argv.extend(
        [
            "--distributed",
            "--distributed-backend",
            "nccl",
            "--fp16",
            "--load", str(model_path),
            "--vocab-file", str(model_path / "vocab.json"),
            "--merge-file", str(model_path / "merges.txt"),
            "--config-path", str(model_path / "deploy.json"),
            "--micro-batch-size", "1",
            "--no-scaled-masked-softmax-fusion",
            "--no-load-rng",
            "--no-load-optim"
        ]
    )
    initialize_megatron(ignore_unknown_args=True)

if __name__ == "__main__":
    if mpu.is_unitialized():
        initialize()
    args = get_args()

    tokenizer = get_tokenizer().tokenizer.tokenizer
    model = get_model(lambda: GPTModel(num_tokentypes=0, parallel_output=False))
    load_checkpoint(model=model, optimizer=None, lr_scheduler=None)
    model.eval()

    inputs = "Hi, how are you doing?"
    input_ids = torch.tensor(tokenizer.encode(inputs).ids).unsqueeze(0)
    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(
        input_ids,
        0,
        reset_position_ids=False,
        reset_attention_mask=False,
        eod_mask_loss=False,
    )
    input_ids = input_ids.cuda()
    position_ids  = position_ids.cuda()
    attention_mask = attention_mask.cuda()
    logits = model(input_ids, position_ids, attention_mask)
    print(logits.norm())
```

* transformers
```python
from transformers_ import GPT2LMHeadModel
from tokenizers import ByteLevelBPETokenizer
import torch

model = GPT2LMHeadModel.from_pretrained(".")
tok = ByteLevelBPETokenizer("vocab.json", "merges.txt", unicode_normalizer="nfkc")

input_ids = torch.tensor(tok.encode("Hi, how are you doing?").ids).unsqueeze(0).cuda()
model = model.cuda()
model.eval()
out = model(input_ids)
print(out.logits.norm())

```

Hi, I also meet the problem about converting Megatron-LM to HuggingFace last week and I want to give my solution, which may help you.
+ The convert code need to be changed as @navatig, where there should be a  lower triangular part of the matrix [change](https://github.com/huggingface/transformers/pull/12007/files)
+ Megatron-LM has four kinds of checkpoint versions, shown as [here](https://github.com/NVIDIA/Megatron-LM/blob/42c1cf4279acea5a554500dcb552211f44cbec45/megatron/checkpointing.py#L209). You have also mentioned the different split ways in your issue. So I suggest you to check the version of your own checkpoint, transform is into version 0 (which is the same as the convert script). This code shows how to transform version 0/1 to version 3, I think you can modify it as you need.

Thank you!
The checkpoint I was using was 3.0 and when I modified the attention matrix, it worked!

@codecaution Thanks a lot for your help! It was very useful!