

TITLE
[new model] `GPTMeg`  (a clone of `GPT2` with a few tiny changes)

COMMENTS
1

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

BODY
# What does this PR do?

This PR tries to address [this issue](https://github.com/bigscience-workshop/Megatron-DeepSpeed/issues/138) by cloning `HF's GPT2` to create `GPTMeg` with a few tiny changes for fp16 adaptation. 

The 3 sources of divergence are:

1. `layer_norm` override to be forced to be done in fp32 - as `MixedFusedLayerNorm` (meg) performs it in fp32 and then casts back to fp16
2. overrides `gelu_fast` to use meg's version which uses `torch.jit.script` which under fp16 returns diverging output (the bwd function was needed to support jit and it's not the issue)
3. `_attn` override to use `torch.baddbmm` instead of `torch.matmul` - the divergence happens due to the alpha factor as it gets applied differently in the 2 ways.


Questions:
- need to decide whether `GPTMeg` is a good name - as BigScience is going to release several more variations of GPT2 - this particular model is however a straight Megatron-LM gpt2 without any changes from BigScience. We discovered that the HF gpt2 model is an almost exact fit wrt to producing the same logits under fp32, but quite a large divergence happens under fp16. So this variant does things differently under fp16.
- I'm not sure whether the tokenizer should be cloned as it's the same? Do we have models that re-use tokenizers from other models and don't provide their own? here GPT2Tok* is just fine
