Add padding for decoder inputs for flax t5 example
Author: changjonathanc
Pull request: yes

# What does this PR do?

Improve TPU training performance by adding padding for decoder inputs in flax t5 training example.


- maintained examples (not research project or legacy): @sgugger, @patil-suraj


Comments:
Thanks for the measuring the speed-up! Seems like it's always a good idea to pad to powers of 2. The label length had a fixed input length previously (118 I think), but it seems that one should always pad to multiples of 2 then. 

Note that with this fix however we need to make sure that no loss is computed on the added padded tokens which in this current form it is. So we need to add a label_mask etc... -> I'll do a more through review on Monday :-) 

You are right, labels should be masked here https://github.com/huggingface/transformers/blob/a76eebfc806b863bf1eb721ba8c49ef9c2f5049f/examples/flax/language-modeling/run_t5_mlm_flax.py#L657

I will update it later.

BTW, the choice of 128 was following the suggestion here: https://cloud.google.com/tpu/docs/troubleshooting#memory-usage

@patrickvonplaten 
Will this cause memory leak? The `train_metric` is `ShardedDeviceArray`, and it's only de-referenced at the end of an epoch. 
https://github.com/huggingface/transformers/blob/a76eebfc806b863bf1eb721ba8c49ef9c2f5049f/examples/flax/language-modeling/run_t5_mlm_flax.py#L720-L721

> @patrickvonplaten
> Will this cause memory leak? The `train_metric` is `ShardedDeviceArray`, and it's only de-referenced at the end of an epoch.
> https://github.com/huggingface/transformers/blob/a76eebfc806b863bf1eb721ba8c49ef9c2f5049f/examples/flax/language-modeling/run_t5_mlm_flax.py#L720-L721

That's a very good point - I'm currently working on updating all those scripts :-)

> > @patrickvonplaten
> > Will this cause memory leak? The `train_metric` is `ShardedDeviceArray`, and it's only de-referenced at the end of an epoch.
> > https://github.com/huggingface/transformers/blob/a76eebfc806b863bf1eb721ba8c49ef9c2f5049f/examples/flax/language-modeling/run_t5_mlm_flax.py#L720-L721
> 
> That's a very good point - I'm currently working on updating all those scripts :-)

One thing to note, convert the metrics to numpy at the end of each step sounds like it would block the process. Did not specifically test it, and I don't know how to move data from tpu to cpu without blocking.

Edit: I found this: https://github.com/google/jax/issues/2851

I think this would work:
```
cpus = jax.devices("cpu") 

firsts = jax.tree_map(lambda x: x[0], train_metric) # take the first shard, might need to modify the code later
# train_metric = jax.device_get(firsts) # this will convert to ordinary numpy array, probably blocking
train_metric = jax.device_put(firsts, cpus[0]) # this will convert to jax device array on cpu, probably non-blocking

```

> > > @patrickvonplaten
> > > Will this cause memory leak? The `train_metric` is `ShardedDeviceArray`, and it's only de-referenced at the end of an epoch.
> > > https://github.com/huggingface/transformers/blob/a76eebfc806b863bf1eb721ba8c49ef9c2f5049f/examples/flax/language-modeling/run_t5_mlm_flax.py#L720-L721
> > 
> > 
> > That's a very good point - I'm currently working on updating all those scripts :-)
> 
> One thing to note, convert the metrics to numpy at the end of each step sounds like it would block the process. Did not specifically test it, and I don't know how to move data from tpu to cpu without blocking.
> 
> Edit: I found this: [google/jax#2851](https://github.com/google/jax/issues/2851)
> 
> I think this would work:
> 
> ```
> cpus = jax.devices("cpu") 
> 
> firsts = jax.tree_map(lambda x: x[0], train_metric) # take the first shard, might need to modify the code later
> # train_metric = jax.device_get(firsts) # this will convert to ordinary numpy array, probably blocking
> train_metric = jax.device_put(firsts, cpus[0]) # this will convert to jax device array on cpu, probably non-blocking
> ```

Did you notice a big speed up by doing so? :-)

Also cc @patil-suraj here.

This looks pretty interesting! Might make the scripts safer and faster

I didn't specifically test it, but I think it shouldn't be an issue if you use dataloader that support pre-fectching, like tfds or pytorch dataloader. And there won't be much of a difference if the cpu time is low.
But this is how I do it anyway:
https://github.com/cccntu/fnet-generation/blob/42cf70ebba26940afc9fbdb21b6576653cb3827c/scripts/run_summarization_flax.py#L811-L832

I played a bit around with `to_cpu_device` and it causes a ~5-10% slow-down which I think is kind of expected given that jax makes use of asynchronus futures. If we force the metrics to be computed after every training step, we can't make good use of device arrays being computed asynchronously. See: https://jax.readthedocs.io/en/latest/async_dispatch.html#async-dispatch

However, this change seems to lead to a huge speed-up: https://github.com/huggingface/transformers/pull/13012

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.