Remove config assumption in Trainer
# What does this PR do?

This PR tries to limit the access to `model.config` in `Trainer` to the minimum so that it works with regular PyTorch modules (as long as they accept dict inputs and return loss first like our models). The most challenging part was the storing/restoring of the `total_flos`, which I moved to the newly created `TrainerState`. It should work as before and be saved along the rest of the training state.