[s2s] fix finetune.py to adjust for #8530 changes

[stas00]

Making the script work again after #8530 change.

As mentioned in #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

[stas00]

@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

[sgugger]

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.

[sgugger]

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.

[sgugger]

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.

[sgugger]

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

[stas00]

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

[stas00]

@ptrblck, can #8612 (comment) 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?

[sgugger]

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

[stas00]

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