[bug] nan validation loss when finetuning BERT sentiment analysis classifier via Trainer API

Enter the chapter number

Chapter 14

Enter the page number

No response

What is the cell's number in the notebook

Cell 120

Enter the environment you are using to run the notebook

Colab

Describe your issue

I'm running chapter 14 in colab, specifically the sections "Task-specific classes" and "The Trainer API" (cells 112-120), and the trainer is reporting a nan validation loss with accuracy between 40-50%, i.e., something completely breaks.

It seemed very odd, since validation metric looks indeed correct and the loss is backed into the Trainer object itself.

Enter what you expected to happen

No response

If you found a workaround, describe it here

Since I was skeptical the issue was with the metric/loss, I checked the float16.

First I tried to remove the conversion, but the T4 of colab was forcing to extremely reduce the batch size making 1 epoch incredibly long. Then I realised that Trainer class offers flags to specify the model weights dtype.

Overall, I resolved the problem by commenting out the dtype=torch.float16 in the instantiation of the pretrained model in cell 112

bert_for_binary_clf = BertForSequenceClassification.from_pretrained(
    "bert-base-uncased", 
     num_labels=2, 
     #dtype=torch.float16      #<---- remove this
).to(device)

and adding type casting using the TrainerArguments options, based on the official documentation in cell 119

train_args = TrainingArguments(
    output_dir="my_imdb_model", num_train_epochs=2,
    per_device_train_batch_size=128, 
    per_device_eval_batch_size=128,
    eval_strategy="epoch", 
    logging_strategy="epoch", 
    save_strategy="epoch",
    load_best_model_at_end=True, 
    metric_for_best_model="accuracy",
    report_to="none",
    fp16=True.       # <--- add this
)

With these two changes I get a validation loss of around 0.3 and (more important) the validation accuracy is around 0.9 (as originally in the notebook).