

TITLE
[WIP] Add Megatron-11B

COMMENTS
8

REACTIONS
+1: 0
-1: 0
laugh: 0
hooray: 0
heart: 3
rocket: 0
eyes: 0

BODY
# What does this PR do?

Fixes #9560

This PR introduces the Megatron model as described in https://github.com/pytorch/fairseq/blob/master/examples/megatron_11b/README.md
This one will probably be fun to test with DeepSpeed, as @stas00 mentioned it's referenced a lot in its docs :smile: 

It's important to mention that there are actually two independent implementations of [Megatron-LM](https://arxiv.org/pdf/1909.08053.pdf):
* The one described in the original paper belongs to NVIDIA (https://github.com/NVIDIA/Megatron-LM), but they released only a 345M checkpoint. It's also based on a rewrite of GPT2 and is not compatible with the current huggingface implementation due to minor changes, like LayerNorm reordering (see https://github.com/NVIDIA/Megatron-LM/issues/37).
* [Fairseq](https://github.com/pytorch/fairseq/blob/master/examples/megatron_11b/README.md), on the other hand, uses its own GPT2 version based on their encoder-decoder framework (with the encoder removed) and it does release the colossal 11B pretrained model.

After some tinkering I realized that fairseq's checkpoint is already pretty compatible with the existing BART port. So, based on that and the fact that NVIDIA doesn't plan on releasing the 3B and 8B checkpoints, **I chose to port only the fairseq version**.

**NOTE:** The original fairseq implementation requires an 8-GPU server to even load the model weights, so I just load the checkpoints manually one by one and merge the model-parallelized tensors into single-model ones.

### How to reproduce the conversion

1. First, find a server with _at least 85GB of RAM_, this model is huge!
2. Next, download and untar the checkpoint:
```
# WARNING: this file is 19GB
wget https://dl.fbaipublicfiles.com/fairseq/models/model_parallel/megatron_11b.tar.gz
tar -xzvf megatron_11b.tar.gz
wget -P ./megatron_11b/ 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/encoder.json'
wget -P ./megatron_11b/ 'https://dl.fbaipublicfiles.com/fairseq/gpt2_bpe/vocab.bpe'
```
3. Run the conversion script
```
python convert_megatron_original_pytorch_checkpoint_to_pytorch.py --fairseq_path /path/to/megatron_11b --pytorch_dump_path /path/to/megatron_hf_dump
```
4. The conversion script will load the model-parallel shards of the checkpoint, group the sharded parameters and concatenate the weights, so that the [fairseq.ModelParallelTransformerLanguageModel](https://github.com/pytorch/fairseq/blob/3b27ed7996b0315f471c795cf9b7dfcc18467cbe/fairseq/model_parallel/models/transformer_lm.py) `state_dict` can be easily loaded into a CPU-compatible [faiseq.TransformerLanguageModel](https://github.com/pytorch/fairseq/blob/3b27ed7996b0315f471c795cf9b7dfcc18467cbe/fairseq/models/transformer_lm.py). The de-parallelisation is based on ParlAI's [conversion script](https://github.com/facebookresearch/ParlAI/blob/abfb771ac4ed2966d6f3ea22c7a38e4ebc9cc0f0/parlai/agents/bart/convert_fairseq_to_parlai.py#L258-L307).
5. Then the script will initialize the huggingface Megatron model and load the converted `state_dict` into it.

### Here's how Megatron differs from the existing BART/MBART implemenations:

1. The most controversial difference, IMO, is the missing encoder, since it's a decoder-only model. For now, I decided to remove the encoder parts inherited from MBART, bit left the encoder-dependent parts in the decoder (e.g. `encoder_hidden_states`, `encoder_attention_mask`) and the cross-attention to simplify the review process on your end. 
2. Megatron uses `SinusoidalPositionalEmbedding` instead of learned ones, so I just yanked those from FSMT :smile: 
3. Megatron does not have a `layernorm_embedding`
4. Minor detail: the `self_attn_layer_norm` is applied before self-attention (like in MBART) instead of after (like in BART).

### Important questions regarding the API:

1. What should be done about the missing encoder? I think the `decoder` variable can be left as is, since it's compatible with the fairseq checkpoint keys, but the `encoder_*` references in the code bother me a lot. We need to somehow strike a balance between `Copied from` and removing the unused parts.
2. I think the position of `self_attn_layer_norm` should be a parameter in the config, similar to `decoder_normalize_before=True` in faiseq. This will close the not-so-obvious difference between BART and MBART.
3. The existence of `layernorm_embedding` can also be parametrized, similar to `layernorm_embedding=False` in fairseq.

### Quick LM test

You can test out the model's capabilities like so (again, you'll probably need _at least 85GB RAM_, there's some weird memory duplication happening somewhere, this should not need more than 50):
```
from transformers import MegatronForCausalLM, MegatronTokenizer, TextGenerationPipeline

tokenizer = MegatronTokenizer.from_pretrained("megatron-11b")
model = MegatronForCausalLM.from_pretrained("anton-l/megatron-11b")

def generate(prompt, max_length=40, num_beams=5, num_return=3):
    input_ids = tokenizer(prompt, return_tensors="pt").input_ids
    outputs = model.generate(
        input_ids=input_ids, num_beams=num_beams, num_return_sequences=num_return, max_length=max_length
    )
    decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True)
    return decoded

print(generate("Before boarding your rocket to Mars, remember to pack these items: "))
```
```
['Before boarding your rocket to Mars, remember to pack these items: 1. A parachute.',
 'Before boarding your rocket to Mars, remember to pack these items: 1. A parachute $100 bill2. A copy of your passport3. A copy of your passport444',
 'Before boarding your rocket to Mars, remember to pack these items: 1. A parachute $1 million dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars dollars']
```

To be honest, I'm not too impressed with its text-generation power. :smile: I guess it's either that the model was too large to train it for enough steps, or I missed something during the conversion. The original implementation does not have a text-generation script (or any non-wikitext results, for that matter), so I'm kinda in the dark here.

## 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.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

@patrickvonplaten, @patil-suraj
