Improving model saving and resuming
Author: bilal2vec
Pull request: no

## 🚀 Feature

The official transformer examples should make it easier to continue training a model that suddenly stopped (e.g. vm gets preempted in the middle of a training run).

To do this, the examples should be updated to save the `optimizer` and `scheduler` states to the `output_dir` as well as the current epoch to disk at the end of each epoch in a `training_state.pt` file. This way, the user could choose to continue training from a previous model checkpoint, but would continue training from the saved epoch and would use the saved tokenizer, optimizer, and scheduler.

## Motivation

If your VM gets preempted in the middle of a training run, you won't be able to properly continue training the model since the scheduler will be reset and the current learning rate will be lost.

## Additional context

If anyone is interested in this, I can implement the feature and start a pull request.

Related issues:

- https://github.com/huggingface/transformers/issues/1925
- https://github.com/huggingface/transformers/issues/839

Comments:
I think it's very useful this feature because, as you highlight, VMs and other variables could stop the (long) training process.
Technically, how could you implement this feature? Surrounding all the training code in a `try/except` statement and when it occurs a particular Exception (which one?) you ask the user whether he/she wants to save till now and after that you saved a file?

It could be necessary to write a method for saving the optimizer state, the scheduler and the tokenizer in a **standardized** way. Reading #1925 and #839, I understand that @thomwolf suggests to use standard PyTorch method for saving and loading e.g. scheduler.

> ##  Feature
> The official transformer examples should make it easier to continue training a model that suddenly stopped (e.g. vm gets preempted in the middle of a training run).
> 
> To do this, the examples should be updated to save the `optimizer` and `scheduler` states to the `output_dir` as well as the current epoch to disk at the end of each epoch in a `training_state.pt` file. This way, the user could choose to continue training from a previous model checkpoint, but would continue training from the saved epoch and would use the saved tokenizer, optimizer, and scheduler.
> 
> ## Motivation
> If your VM gets preempted in the middle of a training run, you won't be able to properly continue training the model since the scheduler will be reset and the current learning rate will be lost.
> 
> ## Additional context
> If anyone is interested in this, I can implement the feature and start a pull request.
> 
> Related issues:
> 
> * #1925
> * #839

This would be very useful indeed. 

I would guess that when the VM gets preempted the process in which your program runs is sent a `SIGTERM` or `SIGKILL` signal from the OS. You would need to catch this signal and act accordingly. Look at the [signal module](https://docs.python.org/3/library/signal.html) in python's standard library.

An elegant and very general solution would be to define a context manager (`with ... do`) in which we would execute the training and that handles all the backup logic on `SIGTERM` or `SIGKILL`.

Do you want to give it a shot and make a Pull Request? You can @ me when you have a first draft and I can have a look and give you feedback.

@rlouf I was thinking to just save the optimizer and scheduler whenever the model is saved. As for resuming, you could just load in the optimizer state, scheduler state, and current epoch from the checkpoint file when passing in `--model_name_or_path`

I've got a basic implementation of this on my [fork](https://github.com/bkkaggle/transformers/tree/saving-and-resuming)


I've seen your changes in source code. I think this is the "easiest" way to handle this feature and I like it. Do you have tested your code with unit testing? I don't see any test suite. Only for being sure that it works as expected.
N.B: I've left some comments under your changes in your repo, please read them.

> @rlouf I was thinking to just save the optimizer and scheduler whenever the model is saved. As for resuming, you could just load in the optimizer state, scheduler state, and current epoch from the checkpoint file when passing in `--model_name_or_path`
> 
> I've got a basic implementation of this on my [fork](https://github.com/bkkaggle/transformers/tree/saving-and-resuming)

@bkkaggle Definitely the easiest solution if you don't mind resuming from the last checkpoint---mine was a bit heavy-duty :)

I also agree with @TheEdoardo93; Can you rebase your branch on the repo's `master` and open a pull request (if you haven't done so already)? I'm happy to have a closer look.


I've updated my branch and submitted a [WIP] [pull request](https://github.com/huggingface/transformers/pull/1987)