gpt2 and t5 parallel modeling
Author: alexorona
Pull request: yes
Labels: name: Model Parallel

# Model Parallelism for T5 and GPT2

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
@LysandreJik 

Comments:
Would it be a good idea to support a less painful way of writing a device_map? This is because as the developer experiments with the mapping, the current method is very inefficient to modify the layer maps.

Perhaps there could be more than one way to do it?

Instead of:
```
                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]}

```

e.g. some ideas for perhaps much simpler ways to create such a map:

* remap the layers as follows:
```
device_map = {
     devices: [0, 3, 4, 5],
     layer_split: [8, 12, 12, 12],
}
```
* a simple string:
```
device_map = {
    devices: [0, 1, 2, 3],
    layer_map: "1-10, 9-21, 22-34, 35-47",
}
```
* a simple string using slice notation
```
device_map = {
    devices: [0, 1, 2, 3],
    layer_slice: "1:10, 9:21, 22:34, 35:47",
}
```

probably several ways can be supported and a wrapper expand them into the explicit version used now based on the keys of the map in the argument.

in either case, changing the map is much easier then...


Not a bad idea. Adding to that: create a mapping utility that's been tested for larger model types like `device_map = get_device_map(machine = 4, model = "gpt2-xl")`. The first device should have fewer layers because it has the embedding and head. This is what I was using for testing:

```
def get_device_map(machine: str, model_name: str) -> dict:
    """Returns a dictionary optimized for distributing a model across
    several devices in a model parallel manner."""
    if machine in ["TeslaV100x4", "p3.8xlarge", 4]:
        device_dict = {
            "gpt2-xl": {
                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],
            },
            "t5-large": {
                0: [0, 1, 2],
                1: [3, 4, 5, 6, 7, 8, 9],
                2: [10, 11, 12, 13, 14, 15, 16],
                3: [17, 18, 19, 20, 21, 22, 23],
            },
            "t5-3b": {
                0: [0, 1, 2],
                1: [3, 4, 5, 6, 7, 8, 9],
                2: [10, 11, 12, 13, 14, 15, 16],
                3: [17, 18, 19, 20, 21, 22, 23],
            },
            "t5-11b": {
                0: [0, 1, 2],
                1: [3, 4, 5, 6, 7, 8, 9],
                2: [10, 11, 12, 13, 14, 15, 16],
                3: [17, 18, 19, 20, 21, 22, 23],
            },
        }
    return device_dict[model_name]
```

This is definitely a goodness to have in the library as well as it will save the developer start up time! I'd even extend it to a specific card size in the argument to `get_device_map`, since the map would be different depending on the size of the cards.

But this won't work for cards of different sizes, e.g. at the moment I have 1x 22GB + 1x 8GB cards. But perhaps this is an odd case and most serious setups have identical cards. I don't know. 


@stas00 Yeah, I think you're right. Could do something simple like `device_map` dictionary should be ranges like this:

```
device_map = {0: range(0, 10),
              1: range(11, 24),
              ...}
```
Simpler than creating a list.

well, it's the same just using python to save on typing ;) this is still awkward a bit as you have to count ;)

here there is less counting: I want you to use devices `[0,1,2,3]` and slice the layers as `[8, 12, 12, 12]` :)

That's why I'm suggesting to support more than one way.

But most likely any of these custom ways can be easily delegated to a helper util, so the end result is the `device_map` as you implemented it. e..g:
```
device_map=device_map_make_by_partition([0,1,2,3], [8, 12, 12, 12])
device_map=device_map_make_by_slice([0,1,2,3], "1:10, 9:21, 22:34, 35:47")
```