[s2s] fix finetune.py to adjust for #8530 changes
Author: stas00
Pull request: yes

Making the script work again after https://github.com/huggingface/transformers/pull/8530 change.

As mentioned in https://github.com/huggingface/transformers/pull/8612 `.logits` doesn't seem to work with apex/PL. No idea why.
So `distillation.py` is probably a problem too since it uses `.logits`. I haven't checked it.

I don't think any tests use apex.
```
RUN_SLOW=1 pyt -sv test_bash_script.py::TestMbartCc25Enro::test_train_mbart_cc25_enro_script
```
tests finetune with fp16, but it runs it in a different way.

@sgugger, @LysandreJik 

Comments:
@sgugger if you want to see why `.logits` isn't working you can try:
```
cd examples/seq2seq
PYTHONPATH="../../src" CUDA_VISIBLE_DEVICES=0 python finetune.py --learning_rate 3e-5 --gpus 1 --do_train --val_check_interval 1 --num_train_epochs 1 --freeze_encoder --freeze_embeds --data_dir cnn_dm --max_target_length 142 --val_max_target_length 142 --train_batch_size 1 --eval_batch_size 1 --gradient_accumulation_steps 1 --model_name_or_path sshleifer/student_cnn_12_6 --tokenizer_name facebook/bart-large --output_dir distilbart-cnn-12-6 --overwrite_output_dir --num_sanity_val_steps 0 --fp16 --eval_beams 1 --amp_backend=apex --n_train 1 --n_val 1 --warmup_steps 1
```

If I `pprint(outputs)` I get a dict.

If you need `cnn_dm`:
```
wget https://cdn-datasets.huggingface.co/summarization/cnn_dm_v2.tgz
tar -xzvf cnn_dm_v2.tgz  # empty lines removed
mv cnn_cln cnn_dm
```

I'll investigate but at first glance it looks like it's PyTorch Lightning that is messing with the new output types, not PyTorch. We will add a caveat in #8530 if that's the case.

So this is not linked to mixed precision directly but purely something in PL. Somewhere when dealing with mixed precision, it looks like there is something happening is the output is an instance of a dict, and they don't respect the class of the dict.

More investigation shows the problem is actually coming from apex, when using `opt_level="O2"` (the default level is fine). For reference, the following code shows apex-converted model lose their output type:
```
from apex import amp
from transformers import BertModel, BertTokenizer

tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer("Hi, my name is Sylvain!", return_tensors="pt")
inputs = {k: v.cuda() for k, v in inputs.items()}

model = BertModel.from_pretrained("bert-base-uncased")
model = model.cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
outputs = model(**inputs)
type(outputs)
```
should return `BaseModelOutputWithPoolingAndCrossAttentions` but returns `dict`.

(The fix is the right one, so while the discussion may continue, I'm merging this PR.)

Thank you for investigating this, @sgugger.

As expected distillation.py fails too with apex / default level 2
```
python distillation.py   --teacher facebook/bart-large-xsum --data_dir xsum   --tokenizer_name facebook/bart-large-xsum   --student_decoder_layers 6 --student_encoder_layers 12   --freeze_encoder --freeze_embeds   --learning_rate=3e-4   --do_train   --do_predict   --fp16   --val_check_interval 0.1 --n_val 1 --eval_beams 1 --length_penalty=0.5   --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100   --model_name_or_path IGNORED   --alpha_hid=3.   --train_batch_size=16 --eval_batch_size=16 --gradient_accumulation_steps=2   --sortish_sampler   --num_train_epochs=6   --warmup_steps 1   --output_dir distilbart_xsum_12_6     --amp_backend=apex     --n_train 1     --gpus 1
[...]
  File "distillation.py", line 157, in _step
    lm_logits = student_outputs.logits
AttributeError: 'dict' object has no attribute 'logits'
```
there are multiple situations like this in this program.

what's the best way to proceed, @sgugger? switch to dict keys for now and report this issue to apex? (which is not being watched - even PRs aren't being merged/attended to).

@ptrblck, can https://github.com/huggingface/transformers/pull/8612#issuecomment-729721261 be fixed in apex, or is it a lost cause (as I noticed apex is not actively supported anymore). Please let me know if we should ask someone else? 

I think we should fix the scripts by accessing elements in the outputs with their keys.

I will do that. Thank you for the feedback, @sgugger