Fix megatron_gpt2 attention block's causal mask
Author: novatig
Pull request: yes

# What does this PR do?
Fixes the conversion between Megatron-LM's GPT2 parameters and transformer's GPT2.
In Megatron-LM the attention mask is implemented differently and is not part of the provided checkpoint.
This PR initializes the mask as in [GPT2](https://github.com/huggingface/transformers/blob/master/src/transformers/models/gpt2/modeling_gpt2.py#L131).

Fixes [#11916](https://github.com/huggingface/transformers/issues/11916) and [#12004](https://github.com/huggingface/transformers/issues/12004).

Regarding [#11916](https://github.com/huggingface/transformers/issues/11916), this PR lowers the perplexity to 20 (cf. 30 for gpt2) without fine-tuning.

Regarding [#12004](https://github.com/huggingface/transformers/issues/12004), the prompt in the issue now is continued coherently.
> How are you doing these days?
> I'm just trying to get through the day. I've been working on a lot of things, but it's hard when your kids come home and they're not here with me anymore because we don't have that connection like before." She said she has had some good times since her son was born in January 2012: "It feels great being able for him [to] be around his dad again," he added as if remembering how much fun their relationship used too! The couple also share two daughters together – Ella Rose (born April 2013) who is now 10 years old; Aviana Grace Marie-Gracee DeSantis Jr., 5th grade daughter from an earlier marriage whom Lohan recently

 ## Who can review?
The PR for the megatron models was reviewed by @LysandreJik
@jdemouth 

Comments:
Hi, I am struggling with the same issue as I detailed in #12004

Did you happen to confirm that once you make this change the two logits produced by Megatron GPT2 and transformers GPT2 are identical?

From what I have experimented, making attention mask as lower triangular matrix did not produce the same logits, nor generated sensible sentences :(

Plus, although I did not check thorougly, [this line](https://github.com/huggingface/transformers/blob/61c506349134db0a0a2fd6fb2eff8e29a2f84e79/src/transformers/models/gpt2/modeling_gpt2.py#L132) seems to cancel out the wrong conversion, you might wanna take a look!

Could you share the generation snippet to reproduce the following result? 
```
How are you doing these days?
I'm just trying to get through the day. I've been working on a lot of things, but it's hard when your kids come home and they're not here with me anymore because we don't have that connection like before." She said she has had some good times since her son was born in January 2012: "It feels great being able for him [to] be around his dad again," he added as if remembering how much fun their relationship used too! The couple also share two daughters together – Ella Rose (born April 2013) who is now 10 years old; Aviana Grace Marie-Gracee DeSantis Jr., 5th grade daughter from an earlier marriage whom Lohan recently
```

Hi,

I ran the scripts that you provided in the issue, the logits with Megatron-LM and transformers are close relative to fp16 precision. There is of course the caveat that the embedding table in transformer's GPT2 has size 50257, while in Megatron-LM is 50304 (therefore we must apply [this](https://github.com/huggingface/transformers/blob/master/src/transformers/models/megatron_gpt2/convert_megatron_gpt2_checkpoint.py#L73) step).
(see attached code).

> Plus, although I did not check thorougly, this line seems to cancel out the wrong conversion, you might wanna take a look!

That is initialization, which we overwrite with the checkpoint

> Could you share the generation snippet to reproduce the following result?

I used the code you provided in the issue!


```
import sys
sys.path.append('/project/Megatron-LM')

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
from transformers import GPT2LMHeadModel
from tokenizers import ByteLevelBPETokenizer


def initialize():
    model_path = "/project/megatron-gpt2-345m"
    sys.argv.extend(
        [
            "--num-layers", "24",
            "--hidden-size", "1024",
            "--num-attention-heads", "16",
            "--seq-length", "1024",
            "--max-position-embeddings", "1024",
            "--tokenizer-type", "GPT2BPETokenizer",
            "--fp16",
            "--load", str(model_path),
            "--vocab-file", str(model_path + "/vocab.json"),
            "--merge-file", str(model_path + "/merges.txt"),
            "--micro-batch-size", "1",
            "--checkpoint-activations",
            "--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
    def model_provider(pre_process=True, post_process=True):
        return GPTModel(num_tokentypes=0, parallel_output=False,
                pre_process=True, post_process=True)
    model = get_model(model_provider)
    load_checkpoint(model=model, optimizer=None, lr_scheduler=None)
    model = model[0]
    model.eval()

    inputs = "Hi, how are you doing?"
    input_ids = torch.tensor(tokenizer.encode(inputs)).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)

    model2 = GPT2LMHeadModel.from_pretrained("/project/megatron-gpt2-345m/")
    tok = ByteLevelBPETokenizer(
            "/project/megatron-gpt2-345m/vocab.json",
            "/project/megatron-gpt2-345m/merges.txt", unicode_normalizer="nfkc")
    
    input_ids = torch.tensor(tok.encode("Hi, how are you doing?").ids).unsqueeze(0).cuda()
    model = model2.cuda()
    model2.eval()
    out = model2(input_ids)

    for j in range(out.logits.shape[1]):
      for i in range(out.logits.shape[2]):
        a, b= out.logits[0,j,i].item(), logits[0,j,i].item()
        assert(abs(a-b) / max(max(abs(a),abs(b)), 0.5) < 0.1)
```




@novatig Thanks for solving my issue！
As for the issue of @hwijeen, where he finds the split in self-attention(QKV and heads) of megatron and huggingface in different order, but you get the close answers. 
I think one possible reason is that there exists three versions of checkpoint in megatron, which needs transform between them, shown as https://github.com/NVIDIA/Megatron-LM/blob/42c1cf4279acea5a554500dcb552211f44cbec45/megatron/checkpointing.py#L209. 
So the version of checkpoint should also be reported when talking about the wrong results.
Or it will be better to modify the convert code for supporting different checkpoint versions.

Hi,

could you perhaps add a `test_modeling_megatron_gpt2.py` file to the tests folder? As this model is not a new model (only a conversion script is required), it could look very similar to [the one of BORT](https://github.com/huggingface/transformers/blob/master/tests/test_modeling_bort.py) (which is also just a conversion script to convert the weights to a BertModel). 

The test should only include an integration test, which checks whether the HuggingFace model outputs the same output tensors (e.g. logits) on the same input data as the original implementation. 

Thanks!

Hi,

As @codecaution pointed out, the problem in my case was that I was working with Megatron checkpoint version 3, whereas the proposed conversion code supports version 0 -- so please ignore the above comments of mine.


Hi,

Many thanks @hwijeen and @codecaution for clarifying the issue!
I updated the the PR, now the conversion script works also for checkpoints generated by recent versions of Megatron-LM.

@NielsRogge @LysandreJik I added a brief integration test.
For simplicity, the reference outputs are not those from Megatron-LM but are from the HuggingFace model, and are close to the ones from Megatron-LM.

Thanks @novatig! 

> @NielsRogge @LysandreJik I added a brief integration test.
For simplicity, the reference outputs are not those from Megatron-LM but are from the HuggingFace model, and are close to the ones from Megatron-LM.

For reference, how close are these to the official Megatron GPT-2 output in terms of magnitude?

> Thanks @novatig!
> 
> > @NielsRogge @LysandreJik I added a brief integration test.
> > For simplicity, the reference outputs are not those from Megatron-LM but are from the HuggingFace model, and are close to the ones from Megatron-LM.
> 
> For reference, how close are these to the official Megatron GPT-2 output in terms of magnitude?

For reference, the test values (taken from a "diagonal" of the returned logits tensor) are:
```
megatron = [4.9492188, -0.2866211, -1.2041016, -4.0351562, -0.5180664, -5.2148438, -1.2412109, -1.8310547, -1.7675781, -4.71875, -0.23901367, -1.0761719, -2.1699219, 0.41235352, -3.8007812, -4.0585938, -2.5292969, -3.3808594, 4.3789062]
hf-gpt2  = [4.9414,    -0.2920,    -1.2148,    -4.0273,    -0.5161,    -5.2109,    -1.2412,    -1.8301,    -1.7734,    -4.7148,  -0.2317,     -1.0811,    -2.1777,    0.4141,     -3.7969,    -4.0586,    -2.5332,    -3.3809,    4.3867]
```
The absolute error is approximately <= 1e-2.

I updated the PR with style changes.
Sorry for all the commits! I did not realize someone had to manually confirm approval for running the tests!

> Hi,
> 
> Many thanks @hwijeen and @codecaution for clarifying the issue!
> I updated the the PR, now the conversion script works also for checkpoints generated by recent versions of Megatron-LM.
> 
> @NielsRogge @LysandreJik I added a brief integration test.
> For simplicity, the reference outputs are not those from Megatron-LM but are from the HuggingFace model, and are close to the ones from Megatron-LM.

Hi! @novatig 
Really thanks for your effort! To make this PR better, I would like to give my own suggestion. 
As Megatron-LM is a popular repo for training large transformer-based model, it will be better to take this into consideration in the convert code, including different mode size and model-parallel.

> Hi! @novatig
> Really thanks for your effort! To make this PR better, I would like to give my own suggestion.
> As Megatron-LM is a popular repo for training large transformer-based model, it will be better to take this into consideration in the convert code, including different mode size and model-parallel.

Yes, that's a good idea!
@codecaution please review my last commit to check if it looks like what you were thinking about.



> > Hi! @novatig
> > Really thanks for your effort! To make this PR better, I would like to give my own suggestion.
> > As Megatron-LM is a popular repo for training large transformer-based model, it will be better to take this into consideration in the convert code, including different mode size and model-parallel.
> 
> Yes, that's a good idea!
> @codecaution please review my last commit to check if it looks like what you were thinking about.

Thanks! Good job!