

TITLE
[WIP] Tensor model parallelism for GPTNeo model. 

COMMENTS
5

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

BODY
## 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