[T5 model parallel] implement input auto-relocation + lots of refactoring/cleanup
Author: stas00
Pull request: yes
Labels: name: Model Parallel

name: WIP

As I commented on in another incarnation of generalizing t5 model parallelism https://github.com/huggingface/transformers/pull/9316 so that it could be easily ported to other models I realized that it's quite unnecessary to try and remap inputs to specific devices where they will be needed in the future ahead of time. Since we have `forward` where we have access to the device of the parameters of that layer - we can completely automate the relocation of inputs to the correct devices just before `forward` is called. So this PR builds upon https://github.com/huggingface/transformers/pull/9316 and:

* [x] creates `@model_parallel_inputs_to_device` decorator used for `forward`, which automatically takes any inputs and puts them on the same device as the parameters of that layer. This allowed a complete removal of most of the `.to()` juggling logic for inputs, which was quite complex and noisy.
* [x] a lot of refactoring to make the MP as little invasive and noisy as possible, and fixing some small issues on the way.

I have tested this with:
``` 
pyt -sv tests/test_modeling_t5.py -k parallel
```
Which I'm not sure covers all bases, but the above tests pass.

@alexorona, please let me know what you think. And if you have real applications besides the great tests you wrote please see if it still works correctly. (It was so awesome having those tests in place! Thank you!) If it looks good and others support this proposal we can then look at doing the same for gpt2 and meanwhile I will look at bart.

@patrickvonplaten, @LysandreJik

Comments:
I don't have an in-depth knowledge of our model parallelism features, so it would be great if @LysandreJik can take a look here as well. 

I think in general, I'm in favor of this PR. However, I'm not sure if a function decorator is better than just having two lines of 
```
if self.is_parallel: 
       # call map to device function
```

in the respective forward function. We've decided against using function decorators in Pytorch at multiple points (gradient checkpointing e.g.), so I'm not convinced it's the better option to do it here. Function decorators do reduce code readability quite a lot IMO.

I'm not sure how your suggestion would work since it needs to be generic, and once inside `forward` the function args are no longer generic. Remember, I'm trying to build a generic functionality that can work out of the box in any `transformers` model and not specific to t5.

The other approach that doesn't need a decorator is to override `self.__call__` via `self.parallelize` to set to a variation of this wrapper.

```
    def parallelize(self, device_map=None):
        $self.__call__ = model_parallel__call__
        [...]

    def deparallelize(self):
        $self.__call__ = nn.Module.__call__
        [...]
```
and:
```
def model_parallel__call__(self, *input, **kwargs):
    
    # get device of any of the params of this layer
    try:
        device = next(self.parameters(recurse=True)).device
    except StopIteration:
        device = None

    if device is not None:

        input = list(input)
        for i, v in enumerate(input):
            if v is not None:
                input[i] = v.to(device)
        input = tuple(input)

        for k in kwargs.keys():
            if kwargs[k] is not None and torch.is_tensor(kwargs[k]):
                kwargs[k] = kwargs[k].to(device)

    return nn.Module.__call__(self, *input, **kwargs)
```
(or could save the original `self.__call__` to be more flexible and to allow for others to override this too)

this in fact is even better since it will have 0 impact on non-MP functionality as this wrapper will be called only under MP.

This is great progress, @stas00! From my perspective, to create a general way of doing model parallelism, we need four things:
* a format for `device_map` that can be used on any model
* `device_map` and `model_parallel` need to be attributes on all models, probably by assigning them to `PreTrainedModel`
* `parallelize()` and `deparallelize()` should be on all models, again probably by assigning them to `PreTrainedModel`
* changes to the forward methods need to be abstracted if at all possible (this is by far the most challenging)

This PR makes a lot of progress, the strongest of which is a potential abstraction/simplification of the changes to the forward method. Not sure if a decorator is the solution. @LysandreJik will have that insight when he's back. I like the suggestion by @patrickvonplaten that it's instead a two line implementation `if self.model_parallel` instead of a decorator. But the BIG thing is if most or all of the code in the forward method can be replaced with with something like:
```
if self.model_parallel:
        hidden_states, attention_mask, position_bias, encoder_hidden_states, encoder_attention_mask, encoder_decoder_position_bias, head_mask, past_key_value = _call__mp(hidden_states, attention_mask, position_bias, encoder_hidden_states, encoder_attention_mask, encoder_decoder_position_bias, head_mask, past_key_value)
```
If we can get that right, it might turn model parallelism from a day or weekend project per model into something that takes a few minutes. Much more scalable and sustainable. 

Supporting non-sequential GPUs could be more trouble than its worth -- not entirely sure on this, it's just my instincts. With the billion + parameter models that we're dealing with -- and all indications are that it's only getting bigger going forward -- it's pretty fair to say that most workflows in enterprise and research will be: 

1. develop locally on a machine with one or maybe two GPUs on a small sized version of a model, and 

2. train a final model on a cloud instance or cluster with multiple identical GPUs. 

Sequential hand-offs between GPUs will be the norm in cases like that, which I think are going to be most of them. 

The other thing I worry about is a challenge with PyTorch 1.5 and 1.6 model parallelism behavior. The seemingly redundant clauses and `set_device` statements are there to prevent PyTorch's inferential logic from moving modules or inputs around after `.to()` assignments have been called. It's very annoying. I don't know if it's fixed in 1.7. You'll notice that output layers like the `lm_head` are always on the first device instead of the last device. A more logical workflow would have the embedding layers on the first device and the output layers on the last device. 

I got that to work just fine in forward passes, but I must've tried 10 different ways to get it to behave in backprop before conceding that for whatever reason PyTorch's quantum device superposition just wouldn't allow it. So the output layers are on the same device as the embedding layers. You'd think that matters for load balance between GPUs, and it does -- for gpt2-xl. But since we're practically limited in most situations to PyTorch's 8 GPU per machine preference (inherited from CUDA), by the time you're at 3 billion parameters the embedding and `lm_head` layers are so small in comparison to the attention blocks that it doesn't matter that they're both on the first device, and a custom `device_map` solves the problem for cases where that matters. The implementation implies that there is an extra hand-off or two of a large tensor between GPUs, but I don't think having a perfectly optimal setup will save even 10% on training time. Happy to be proven wrong on this though. What will save a TON of time and $$$ though is deepspeed integration.

I got t5-11b with 1024 tokens to train quickly on the new p4 instance AWS released last month with its epic 320 GB of GPU memory so I was like "ok fine whatever... that's pretty good". 

That's awesome, @alexorona. Do continue to share your insights from the frontier!

Let's wait for @LysandreJik to come back to plan ahead and meanwhile I will experiment with Bart.

> The seemingly redundant clauses and set_device statements are there to prevent PyTorch's inferential logic from moving modules or inputs around after .to() assignments have been called. It's very annoying. I don't know if it's fixed in 1.7. 

Oh, so glad you flagged that. Would it be enough to run the existing parallel tests with pt-1.5, pt-1.6 to detect these failures?

I'm developing on pt-nightly since rtx-30* work only there (well 1.7.1 should be usable too, but mainly waiting for cuda-11.2 support in pytorch, which is again pt-nightly - won't be in 1.7.x). But it means I can't use it with older pt versions.

But since we have to support pt-1.4, I will then put `set_device` back as you had them originally. But this time let's add specific comments why there are there, otherwise someone like myself will think they are some left-overs from earlier experimentation and swipe them away.

Actually, I think we should have a design document where we explain why this or that is done. Rather than make a lot of noise in the model files. A developer-oriented doc. 

The `set_device` was just one thing, right? Or have I naively nuked any other essentials?

Thanks again!

@alexorona, one more question. If pt-1.7+ removes the need for jumping through hoops, as you're suggesting older versions have all kinds of issues, perhaps it'd be a reasonable approach to make MP in `transformers` require pt-1.7?

If and when you get some time could you please test if what wasn't working in pt < 1.7 works in pt-1.7? And if not - perhaps we need to file some Issues with pytorch if there are bugs to be solved.

Thank you.

@stas00 Will try to do so, but in the middle of moving so I don't think I'll get to this until the end of January at the soonest. The team would have to make the call about only support model parallelism for PyTorch >= 1.7.0 if it won't work on earlier versions. I would be very tempted to support that idea, but don't have enough usage information to know what the impact would be.

I guess once everybody is back next week we can start having some discussion with the HF team.

Have an easy move!

@stas00 Yeah, should be able to get some input when everyone is back. In the meantime, I'm still not sure on the final form of the `device_map`. There are two issues left to work out: 

1. Some models don't have decoder architectures
2. No ability to map embeddings and output layers (always on first device), which _might_ be just fine. I think most output layers and embeddings are going to be comparatively small to attention blocks going forward, but we should confirm that. We are allowing people to create custom a `device_map` that should enable them to get around any potential situations where the first device is becoming overloaded.

To confirm, this looks good for decoder architectures:
```
device_map = {
	encoder: {
			0: [0, 1, 2, 3, 4, 5],
			1: [6, 7, 8, 9, 10, 11]
			},
	decoder: {
			2: [0, 1, 2],
			3: [3, 4, 5]
			}
}
```

Maybe we use the keys to map to the attribute? In gpt2, `self.h` contains the attention blocks, so:
```
device_map = {
	h: {
		0: [0, 1, 2, 3, 4, 5],
		1: [6, 7, 8, 9, 10, 11]
		}
}
```
In trying to generalize `parallelize()`, we still need access to list of all modules. For example, in `GPT2LMHeadModel`, we would need to know: `self.lm_head`, `self.transformer.h`, `self.transformer.wte`, `self.transformer.wpe` and `self.transformer.ln_f`. 

I haven't looked into gpt2, yet. t5 and bart are very similar structure-wise. We probably need to map out all the different archs `transformers` has and then generalize.

What is emerging so far is that the device map might have various keys, none required, and each model architecture will have:

1. its required keys
2. its own default map generator - so that the user doesn't have to provide one and overtime it can be improved to have smarts to create a balanced map based on the "insider" information.

So if some architectures need to explicitly manage the mapping of non-block/layers, rather than just assigning them by default on the "`main_device`", because they are significantly big, they could do that too. Otherwise, leave the `main_device` to all the "smallish-fish" and use the other devices for "the big fish" if that makes sense. The main advantage of this "lazy" approach is that there is less device-hopping and less code needed to match the hopping.

Yes, that's right. 

So it turns out the `self._modules` attribute has all of the modules. To move `parallelize()` to `PreTrainedModel`, I think all we need is a per-model `module_map` object to map between the `device_map` and the model placements. With a little work, we might be able to reduce making a model parallel to:

1. Adding a few lines of code in the forward method per your work
2. Modifying the validation function to check for errors in a custom `device_map`
3. Creating a `module_map` dictionary for that model and adding it to the `get_module_map()` function 

We can embed special placement rules where non-attention block modules need to be on the same device as another module by creating a tuple in `module_map['dependent_modules']`:

```
# Device map for GPT2LMHead. T5 would have 'encoder', 'decoder' as keys instead of 'h' and validate_device_map would
# check to see if the device_map has the right keys.
device_map = {
	'h': {
			0: [0, 1, 3, 4],
			1: [5, 6, 7, 8],
			2: [9, 10, 11, 12]
	}
}


class PreTrainedModel():
...
	# Probably use get_model_map(), but just to make it simple:
	self.module_map = {
		'h': self._modules['transformer'].h,
		'embeddings': [
					self._modules['transformer'].wte,
					self._modules['transformer'].wpe
					],
		'dependent_modules': [
					(
						self._modules['transformer'].ln_f], 
						model._modules['transformer'].h[-1],
					),
					(
						self._modules['lm_head'],
						self._modules['transformer'].wte
					)
				]
	}

def parallelize(self, device_map = None):
	
	self.device_map = device_map

	# validate_device_map extended to check for valid keys for model
	
	...

	# Set all embeddings to first device
	if 'embeddings' in self.module_map:
		for layer in self.module_map['embeddings'].items():
			layer.to(self.first_device)


	# Assign attention blocks to the appropriate device.
	for module_group, group_map in self.device_map.items():
	    for device, layers in group_map.items():
	        for layer in layers:
	            self.module_map[module_group][layer].block_parallelize(f"cuda:{device}")


	# Some modules should always be on the same device as another module. We can express 
	# this as a tuple pair where tuple[0] needs to be on tuple[1]
	if 'dependent_modules' in self.module_map:
		 for i in self.module_map['dependent_modules']:
		 	i[0].to(i[1].device)
```


All, awesome suggestions that should be looked at next once the current work has been merged.

I'm going to wait implementing anything new, since there are already too many partial PRs that need to be carefully merged and rebased and once that is done we can do another round of generalization integrating your suggestions.

So if there is no objection, I will merge this one, and then start integrating with https://github.com/huggingface/transformers/pull/9384, which is ahead functionality-wise - so I want to sync the two, switching t5 to the improved version of MP backend. I will implement the suggestions in that new PR.

As we have discovered the original PR didn't make t5 work with trainer. I have just fixed that in the last commit here, bringing some goodies over from the Bart MP PR.

So this now works:

```
export BS=20; rm -r output_dir; CUDA_VISIBLE_DEVICES=0,1 PYTHONPATH=../../src USE_TF=0  ./finetune_trainer.py --model_name_or_path t5-small --output_dir output_dir --adam_eps 1e-06 --data_dir wmt_en_ro --do_eval --do_predict --do_train --evaluation_strategy=steps --freeze_embeds --label_smoothing 0.1 --learning_rate 3e-5 --logging_first_step --logging_steps 1000 --max_source_length 128 --max_target_length 128 --num_train_epochs 1 --overwrite_output_dir --per_device_eval_batch_size $BS --per_device_train_batch_size $BS --predict_with_generate --eval_steps 25000  --sortish_sampler --task translation_en_to_ro --test_max_target_length 128 --val_max_target_length 128 --warmup_steps 500 --n_train 200 --n_val 200 --n_test 200 --fp16 --save_steps 2 --model_parallel
```

But! while it's fine in the training stage, it's 10x slower on eval than w/o `--model_parallel`

Hello @stas00 kudos to all the hard work you do, especially around continuing the ambitious work around supporting parallelism. 

Interested in doing some inference with the t5-11b model variant.
Can you provide some insights on how many gpus would be needed to achieve that?

I tried this branch with 8xV100 (16gb) on GCE. 
All good while I created the model and called parallelize, but got a out of memory error on inference step when moving inputs to the first gpu device.

Let me know if I have a wrong mental model about achieving this. Thanks again!




Thank you for the kind words, @kznmft!

Please have a look at https://github.com/huggingface/transformers/pull/9765
which implements a very inefficient in my opinion but nevertheless working pipeline parallelism on t5, which should be superior to this naive implementation, speed-wise but it's not quite there yet. Please read the first post carefully for all the details. and you can see the follow up comments with the experiments that have been done. So 4x40gb A100s gpus weren't enough for t5-11b in initial experiments. But 5-6 of those probably should be enough.

I finally got access to a machine with 4 gpus just now, so I'm going to start looking at implementing 2D parallelism - using Pipeline with DeepSpeed ZeRO-DP, so I will post news once I get something working.

Subscribe to watch https://github.com/huggingface/transformers/pull/9765 and I most likely will update that PR with new info or a link to a new PR once I have something working.

-----

> I tried this branch with 8xV100 (16gb) on GCE.
> All good while I created the model and called parallelize, but got a out of memory error on inference step when moving inputs to the first gpu device.

But you're not telling me the device map you were using. You need to spread out the layers over the 8 gpus, have you done it? unless you were relying on the default map which should spread things out.

The problem is that it doesn't take into an account that gpu 0 is always overtaxed, so I'd always try a few layers less on the first gpu 0. And then watch nvidia-smi (and later we will have better tools) to see that you get each GPU getting a somewhat equal memory allocation.

But if 4x40 couldn't fit it, I doubt that 8x16 will.

Remember in t5-11b you have 45GB of params, plus optimizer states plus gradients.

Also probably need to try to use a more lean optimizer, say Adam instead of AdamW which needs more memory.





too long. closing.