[examples] SummarizationModule improvements

[sshleifer]

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, 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.

[LysandreJik]

This all looks very cool, looking forward to using it!

[sshleifer]

Merging now. Happy to address post-merge comments!