

TITLE
[model parallelism] Bart goes parallel

COMMENTS
24

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

BODY
This PR implements model parallelism (MP) in Bart.

This is the latest incarnation of generalization of the MP in `transformers`, based on @alexorona's original work. I have done some of it already in https://github.com/huggingface/transformers/pull/9323 and this PR builds upon the other one. It's slightly complex what to merge when, but this PR is independent and can be merged on its own.

For reviewers I propose to read things in this order:
1. https://github.com/huggingface/transformers/pull/9316
2. https://github.com/huggingface/transformers/pull/9323
3. this PR
4. Additional important design discussions https://github.com/huggingface/transformers/issues/8771

If all is in agreement, I propose:

1. &#9744; merging this PR first, 
2. &#9744; then I'll backport the new code from this PR to https://github.com/huggingface/transformers/pull/9323 and we merge that. 
3. &#9744; then we handle gpt2, which I haven't touched yet. Perhaps @alexorona could help there if his time permits or one of us.
4. &#9744; complete Bart's other heads (can be item 3) and `deparallelize` - the latter is not really needed in practice so will handle those when dust around design settles.
5. &#9744; add Bart to trainer's supported for `--model_parallel` flags
6. &#9744; write tests for `model_parallel_utils.py`
7. &#9744; meanwhile we can polish the concept of device maps which will require a review of all architectures `transformers` has implemented.

Actually first we need to merge smaller bits:
1. https://github.com/huggingface/transformers/pull/9347
2. https://github.com/huggingface/transformers/pull/9386

---------

So this PR:

* [x] Implements MP in Bart based on discussions in all of the threads/PRs listed above. Only `BartForConditionalGeneration` at the moment while we are sorting out the API. But the bulk of the work is done, since `BartModel` has all in place.

* [x] switches to the concept of `main_device` rather than `(first|last)_device` so the first device of encoder becomes the main_device and almost everything happens there (`embeddings`, `lm_head`, etc), and other devices are used exclusively for encoder and decoder purposes.

* [x] switches to a more explicit `device_map` that can support non-symmetrical models (different number of layers in encoder and decoder).  It can also handle different types of maps. See the demo at the end this post for details.

* [x] further improves the magical `to()` functions that can operate on any type of variable except opaque objects. Can be used to put the inputs on the correct devices either automatically via a `forward` decorator or explicitly inside `forward`. We could use either or both.

* [x] adds a bunch of debug functions that make it easy to trace device IDs of variables, params and whole layers.

* [x] further improves the device map validation function

* [x] improves tests

* [x] needs to remove  apex.normalization.FusedLayerNorm as it's buggy under MP (corrupts data) per https://github.com/huggingface/transformers/issues/9377 a dedicated to removal PR is https://github.com/huggingface/transformers/pull/9386

Here is a quick demo (you will need 2 gpus to run it):

```
from transformers import BartTokenizer, BartForConditionalGeneration, BartConfig

#mname = "sshleifer/tinier_bart"
mname = "sshleifer/distilbart-xsum-6-6"

model = BartForConditionalGeneration.from_pretrained(mname)
tokenizer = BartTokenizer.from_pretrained(mname)

sentences = ["I'm sitting here in a boring room. It's just another rainy Sunday afternoon. I'm wasting my time I got nothing to do. I'm hanging around I'm waiting for you. But nothing ever happens. And I wonder."]
inputs = tokenizer(sentences, max_length=1024, return_tensors='pt', truncation="longest_first")

device_maps_flat = {
    "sshleifer/tinier_bart": {
        "encoder": {0: [0, 1] },
        "decoder": {1: [0] },
    },
    "sshleifer/distilbart-xsum-6-6": {
        "encoder": {0: [0, 1, 2, 3, 4, 5] },
        "decoder": {1: [0, 1, 2, 3, 4, 5] },
    },
}

device_maps_split = {
    "sshleifer/tinier_bart": {
        "encoder": {0: [0],
                    1: [1],
                    },
        "decoder": {1: [0] },
    },
    "sshleifer/distilbart-xsum-6-6": {
        "encoder": {0: [0, 1, 2],
                    1: [3, 4, 5],
                    },
        "decoder": {0: [0, 1, 2],
                    1: [3, 4, 5],
                    },
    },
}

# 3 different ways (2 different device maps and 1 autogenerated device map)
model.parallelize() # autogenerated
#model.parallelize(device_maps_flat[mname])
#model.parallelize(device_maps_split[mname])

inputs = inputs.to("cuda:0")
# Generate Summary
summary_ids = model.generate(inputs['input_ids'], num_beams=4, max_length=25, early_stopping=True)
print([tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summary_ids])
# prints: [" I'm sitting in a room where I'm waiting for something to happen."]
```

You can see from the demo, that when calling `model.parallelize` you can skip the `device_map` arg altogether and the model will generate the right one. Or you can provide one that:
1. gives some gpus exclusively to encoder and others to decoder
2. splits the model horizontally so that the encoder uses all gpus so so the decoder
the model transparently handles all the remappings

Note, the user still needs to put the data on the `main_device`, so perhaps that will eventually become not hardcoded via:
```
# inputs = inputs.to("cuda:0")
inputs = inputs.to(model.main_device)
```

As we have been discussing elsewhere the device map format is unstable yet. So I propose we document it as unstable yet, but the users can rely on the autogenerated device map which requires no input from the user (i.e. calling `model.parallelize() ) - if it changes it'll happen transparently for the user.

Also note that in situations of Trainer-based scripts, like `finetune_trainer.py`, the user has no way to supply such device map at the moment so in effect the model generates the map on the fly as in the above para.

Fixes: #8344

@LysandreJik, @patrickvonplaten, @sgugger, @alexorona 