Added gpt2 model parallelism
Author: alexorona
Pull request: yes
Labels: name: Model Parallel

# Model Parallelism for GPT2LMHead
Addresses [issue 7526](https://github.com/huggingface/transformers/issues/7526)

Adds two new methods to `GPT2LMHead` and the `GPT2Model` classes to enable you to generate and fine-tune models using model parallelism. This feature is most applicable for `gpt2-large` and `gpt2-xl`. Minor modifications are made to the `TrainingArguments` and `Trainer` classes to avoid conflicting data parallelism behavior and related batch_size increases which would negate model parallelism. Note that nearly 64GB of GPU (4 Tesla v100s)  are needed to fine-tune `gpt2-xl` @ 1024 tokens.

It is critically important to provide users the ability to specify where to put the blocks of a model because the GPU sizes and numbers are likely to be very diverse. This is done with a dictionary called `device_map`. I am planning on providing some examples and guidelines for the p3, p2 and g3 AWS instances. 

Model parallelism has to be baked into the model class itself. Currently working on the T5 model. From my calculations the 11B model cannot fit on the largest p3 instance that I have access to (8 Tesla v100 GPUs). The 3B model can.

The methods are:
- `parallelize`, which will distribute the attention blocks of the model across several devices according to a device map
- `deparallelize`, which will move the model back to cpu

# Example

```
model = GPT2LMHeadModel.from_pretrained('gpt2-xl')
device_map = {0: [0, 1, 2, 3, 4, 5, 6, 7, 8],
                              1: [9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21],
                              2: [22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34],
                              3: [35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47]}
model.parallelize(device_map) # Distributes the model's attention blocks across several devices
model.deparallelize() # Puts the model back on cpu and calls torch.cuda.empty_cache() to liberate GPU memory
```

## Reviewers
`TrainingArguments`: @sgugger 
- Added a new parameter `model_parallel` and attribute `self.model_parallel` to control model parallelism behavior
- Slightly modified the `train_batch_size` and `eval_batch_size` calculations to avoid automatically increasing the batch size if `self.model_parallel` (automatically increasing the batch size defeats the purpose of model parallelism because you won't be able to train a larger model if the batch_size increases proportionally to the number of devices)

`Trainer`: @sgugger
- Minor changes controlled by new `args.model_parallel `attribute

`GPT2LMHead`: @patrickvonplaten 
- Adds parallelize and deparallelize methods
- Adds new `self.model_parallel`  and `self.device_map` attributes
- Changes forward behavior when `self.model_parallel == True` to ensure tensors are on the right device


Comments:
Very cool implementation that would interest a few of our team, pinging them here :)

Indeed, adding a test would be nice! Will trigger the multi-gpu slow tests once the test is added.

@LysandreJik @patrickvonplaten @sgugger Glad you like it! Yes, the basic form of the implementation works on T5. I've been working on that too. Testing solutions to the items you brought up locally and should make the improvements to the PR in the next day or two.

Hi @alexorona! Before we merge this, having a test would really be necessary. Have you had any luck in implementing such a test?

If you're lacking time, we can also take over from here.

Hi, I've successfully been able to start training gpt2-xl on multiple gpu's using the model parallelism code from this pull request, but I'm running into an issue when restarting the training process from a checkpoint model. It seems that in this case when reloading the model, only the memory of the first GPU keeps increasing until it reaches an Out Of Memory error, instead of spreading over all GPUs like it did when training from scratch. Is it possible that reloading from checkpoint triggers a different code path that has been overlooked until now?

(edit)
I seem to have found an issue and potential fix. Line 609 in Trainer.train() loads the optimizer saved in the checkpoint. This function has a _map_location_ parameter which seems to force the optimizer to load fully onto _self.args.device_ which I'm guessing would be the first GPU. 
`self.optimizer.load_state_dict(torch.load(os.path.join(model_path, "optimizer.pt"), map_location=self.args.device))`
Removing the _map_location_ parameter makes the function properly put back all the loaded parameters onto the correct devices. Maybe there is a better way of handling this but at least it is an indication.

@LysandreJik Just haven't had the time to do that and it might be awhile before I can get around to writing tests. However, I have been able to implement the same thing on T5 and it's working. Fine-tuned a 3B model this weekend. Maybe I can add that to this branch tonight and the team can handle the tests? 

@MichielRuelens Did you try loading to CPU and then calling model.parallelize() ? 

@LysandreJik I worked on t5 last weekend. Added that. Happy to explain things like why `get_device_map `should probably take a list of devices rather than a number. Can you help get this across the finish line?

Very cool, thanks! Will take a look at it and open a PR on your branch.

Great, with the tests PR merged this seems close to merge! Could you rebase & run the code quality tools (`make fixup` or `make style && make quality`) so that the test suite passes?

Also @patrickvonplaten and could you check the tests?

Awesome, LGTM!

oof that's a tough rebase. I don't think we'll be able to merge that. Could you close the PR and open a new one, so that we can see the diff? I don't think you need to do anything on the branch, just close the PR and open a new one.

I'm researching doing the same for BART (https://github.com/huggingface/transformers/issues/8344) and stumbled upon the open source library DeepSpeed/ZeRO:
https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/

I'm yet to experiment with deepspeed and fairscale, but I thought I'd ask whether you have already done this and decided that implementing from scratch is better. 

What are the pros/cons for using an in-house implementation as compared to using an external library, other than the obvious potential bugs and issues of using external libs and not having control over those?

If this has been discussed already please kindly send me to that issue/page? Thank you!

Whatever the outcome, we can also qualititatively compare the results of this PR once it's done with doing the same via deepspeed and/or fairscale.


Rebased with transformers 4.0.0 and moved PR to [here](https://github.com/huggingface/transformers/pull/8696).

> I'm researching doing the same for BART (#8344) and stumbled upon the open source library DeepSpeed/ZeRO:
> https://www.microsoft.com/en-us/research/blog/zero-deepspeed-new-system-optimizations-enable-training-models-with-over-100-billion-parameters/

Great point, @stas00 ! From the description, DeepSpeed is an optimization on top of data parallelism and model parallelism.  I think you've identified the next step! My reading of DeepSpeed this that one still has to implement data parallelism or model parallelism, but DeepSpeed will reduce the GPU memory footprint. Practically speaking, this could open the door to even larger models. During the development of the [final model parallel PR](https://github.com/huggingface/transformers/pull/8696), I ran into hard limits on AWS GPU instance with t5-11b and t5-11b in terms of the number of tokens you can train. You'll need the largest AWS instances to train 512 tokens on t5-3b. For t5-11b, you're restricted even more. Note that I haven't tried this with apex, so it might be possible to squeeze out a little more from the current implementation.

Thanks for the follow up, @alexorona. From what I read deepspeed implements parallelism too, but I'm still in the process of studying - didn't get to play with it yet. I've started with fairscale first.

What about [`fairscale`](https://github.com/facebookresearch/fairscale) then? It implements the model parallelism w/o needing to make any changes to  the model. All that code is abstracted into the trainer-level calls. 

Now that you added T5+gpt2 model parallelism `transformers` needs to add a separate code for each model architecture. Why not follow `fairscale`-style and do it outside the model?

Similarly to what we've done with the `tie_weights` methods which ties weights thanks to the `get_input_embeddings()` and `get_output_embeddings()` methods, we could probably have a model-agnostic way of enabling parallelization using a `get_layers()` method, like [this](https://github.com/huggingface/transformers/blob/gradient-checkpointing-v2/src/transformers/modeling_utils.py#L669-L682) for example. 

This would allow to have the same code for each model architecture, provided the `get_layers()` utility is correctly implemented.

Would need to look into how `deepspeed` is implemented. My reading is that it supports parallelism, but still requires you to go through the process of defining how parallelism will work in the model definition (which modules on which devices). 

Not sure on `fairscale`. It may or may not be a simplification. There were quirks in implementing model parallelism on gpt2 and t5 that have to do with the pytorch graph:
- The LM head has to be loaded on the same device as the embedding layer
- Tensors (not just layers) have to be shifted to the appropriate device during training

Initially, I went to `eisen` to see if it could handle model parallelism, but it couldn't deal with transformer models. If `fairscale` can abstract that process, that would be great.

I hear you, @alexorona, that you're saying that each model may have its special needs. So we need to see if perhaps this can be done via some sort of callbacks and add the callbacks points as we discover new needs.

As I was planning to do the same for BART, perhaps I should replicate your implementation for Bart and then based on 3 different models we can then see how we can make this model-agnostic. What do you think?

Perhaps we should continue this discussion in a dedicated issue where we discuss model_parallelism for all of `transformers`, including considering using or forking existing implemenations.

I think those are both great ideas, @stas00. Should give us a better understanding of how to move forward with this. Created [this issue](https://github.com/huggingface/transformers/issues/8771) to continue conversation.