[WIP] Tensor model parallelism for GPTNeo model. 
Author: hyunwoongko
Pull request: yes
Labels: name: Model Parallel

name: Tensor Parallel

name: WIP

## Changes
1. Add `parallelization_utils.py`
2. Add parallel mpdules like row & column parallel linear and vocab parallel embedding and cross entropy to `modeling_utils.py`
3. Add `vocab_parallel_embedding` and `make_vocab_size_divisible_by` attribute to `GPTNeoConfig`
4. Add `GPTNeoParallelismPolicy` in `configuration_gpt_neo.py`
5. Now, `GPTNeoPreTrainedModel` inherits `ParallelizationMixin` class. if a model inherits `ParallelizationMixin`, a method named `from_pretrained_with_parallel()`and `from_config_with_parallel()` will be available. This API is motivated from [parallelformers' API](https://github.com/tunib-ai/parallelformers#2-put-the-model-in-the-parallelize-function)
6. Add `mpu` or `process_group` into trainer's DDP wrapper.

---

## Model parallelism API

- Note this API can be used for both training and inference.
- We can branch `from_pretrained_with_parallel()` into the methods such as `from_pretrained()` + `parallelize()`, but it is designed to create and parallelize the model at the same time for the implementation of pipeline parallelism later. (we should match API of transformers-friendly and megatron-friendly)
- When extending this API to other models in the future, we also need to decide how to dispose the existing naive `parallelize()` method in GPT2 and T5. So I designed the names not to overlap, so that naive method can be deprecated slowly in the future.

```python
'''
deepspeed --num_gpus=4 test_model_parallel_inference.py

or

python -m torch.distributed.launch --nproc_per_node=4 test_model_parallel_inference.py 
'''

from transformers import GPTNeoForCausalLM, GPT2Tokenizer

# FP32
model = GPTNeoForCausalLM.from_pretrained_with_parallel(
    "EleutherAI/gpt-neo-1.3B", tensor_model_parallel_size=4
)

# FP16 (default=False)
model = GPTNeoForCausalLM.from_pretrained_with_parallel(
    "EleutherAI/gpt-neo-1.3B", tensor_model_parallel_size=4, fp16=True,
)

# Use Vocab parallel embedding (default=False)
model = GPTNeoForCausalLM.from_pretrained_with_parallel(
    "EleutherAI/gpt-neo-1.3B", tensor_model_parallel_size=4, vocab_parallel_embedding=True,
)

# For scratch training
model = GPTNeoForCausalLM.from_config_with_parallel(
    YOUR_CONFIG_OBJECT, tensor_model_parallel_size=4
)
```
---

## TODO
- We should add test codes
  - FP32 inference test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM
  - FP16 inference test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM
  - Vocab parallel embedding inference test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM
  - FP32 training Test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM
  - MixedPrecision training test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM
  - Vocab parallel embedding training test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM
  - 2D parallelsim (DP + TP) with Trainer API test
    - [ ] ForSequenceClassification
    - [ ] ForCausalLM 
- [ ] We should consider how to provide checkpoint merging tools
---

## Reviewers

@stas00

Comments:
Note current PR is about transformers-friendly model parallelism, but some codes are designed for both transformers-friendly and megatron-friendly model parallelism. These will be used later when implementing `ParallelGPT2` and `ParallelBert` described [here](https://github.com/huggingface/transformers/issues/13690). I uploaded them because they are the codes must be used in both parallelization methods (transfo-friendly, mega-friendly).

Note I didn't upload the code needed only for megatron-friendly. All of the codes uploaded now are also necessary for transformers-friendly.

@stas00 I'm not familiar with the transformers test suite. Could I ask you to write the test code? Usage is very simple. Just call the `from_pretrained_with_parallel()` method and run with `torch.distributed.launch --nproc_per_node=n` or `deepspeed --num_gpus=n`. Both training and inference require testing.

I think it would be good to divide the roles in our collaboration like this. 
- @hyunwoongko: implement the most of code and test it internally and then PR it. plus I should describe the API so that you can test the code.
- @stas00 or other HF engineers: design test cases that fits the huggingface test suite.

The nice thing about this implementation is that we have an MPU. When using DDP in Trainer, if we put `get_data_parallel_group()` together in MPU, 2D parallel training will be possible (TP + DP). So, I put it in most places where `mpu` or `process_group` is needed. However, I am not familiar with the internal structure of the `Trainer`, so I would like other huggingface engineers to help me with this. 

+) I think we should prevent use of non-distributed data parallel. is it right?

Training with vocab parallel embedding has problem. vocab parallel cross entropy is necessary when language model training (clm) So I added vocab parallel cross entropy loss function.

I will reopen the PR. thanks.