[examples] SummarizationModule improvements
Author: sshleifer
Pull request: yes
Labels: name: seq2seq

name: Examples

This PR makes the SummarizationTrainer much more usable, and when improvements are not unique to summarization, they are implemented in `lightning_base.py` instead.

- **Checkpointing** Before this PR, the code saves 5GB of PL checkpoints per epoch, now SummarizationTrainer saves the best checkpoint based on ROUGE 2 score, and also saves it in huggingface `save_pretrained` format using the `on_save_checkpoint`. This will help resolve lots of confusion in various issues about how to load the pl checkpoints.

The current summarization code can only accept bs=1 and takes 24h to run 1 epoch on CNN DM. With the following changes, you can train much faster, if you wish. The docs suggested that larger batch sizes were possible with default params, which is fixed.

### Changes to Allow Faster Summarization Training
*these are all optional and turned off by default*


1) freezing: before this PR, it was basically only possible to finetune with batchsize 2-4 on a 16GB system. With `--freeze_embeds` and `--freeze_encoder`, you can get batch size MUCH higher, towards 32. I've seen strong results with these options. 

2) On CNNDM and XSUM the datasets are 200K examples, and epochs are very long. For this reason it is preferable to run validation (and get a rouge score) more frequently, but with previous params each `validation_step` took 1hr. By passing `--n_val=1000 --val_check_interval=0.25`, you can run validation 4x per epoch and it only takes 3 minutes. I also allows the config's beam search parameters to be used, rather than hardcoding faster but lower scoring ones. 

3) `{train|val|test}_max_target_length`: I have found it preferable to truncate train summaries to 56 for XSUM and CNNDM respectively, but doing this for val/test artificially inflates rouge scores. So these clargs are separated.


Changes to `lightning_base`
- Number of trainable parameters and total parameters are logged by default.
- All possible `pl.Trainer` clargs are passed through `add_generic_args` (Inspired by @nateraw)


### WandbLogger
- `--logger wandb` will instantiate a default wandb logger.
- `--logger wandb_shared` will post results to [here](https://app.wandb.ai/sshleifer/hf_summarization/table?workspace=user-), so that the community can compare hyperparameter settings empirically.
- the default logger is still tensorboard logger because it doesn't require making an account.

### Distillation
- `SummarizationDistiller` and `T5SummarizationDistiller` are checked in. This code was sent to me by a researcher who wishes to remain anonymous. DM to discuss.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=h1) Report
> Merging [#4951](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/c852036b4abca2c20e1adf92eda48472a7d84ef0&el=desc) will **increase** coverage by `0.00%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/4951/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master    #4951   +/-   ##
=======================================
  Coverage   77.41%   77.42%           
=======================================
  Files         130      130           
  Lines       22023    22023           
=======================================
+ Hits        17050    17051    +1     
+ Misses       4973     4972    -1     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_bart.py](https://codecov.io/gh/huggingface/transformers/pull/4951/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iYXJ0LnB5) | `96.26% <ø> (ø)` | |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4951/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `87.42% <0.00%> (ø)` | |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/4951/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `75.40% <0.00%> (+0.40%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=footer). Last update [c852036...b7e1d5e](https://codecov.io/gh/huggingface/transformers/pull/4951?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Merging now. Happy to address post-merge comments!