Saving and resuming
Author: bilal2vec
Pull request: yes

Here's my basic implementation of the saving and resuming improvements discussed in #1960. So far, I've only modified the `run_lm_finetuning` example, but if my changes are approved I can update the rest of the examples as well.

There are three main changes:

1. The example now saves the optimizer, scheduler, and tokenizer every `save_steps` iterations.
2. The example now checks whether training is being continued from a checkpoint, and if so, looks for a saved optimizer and scheduler and loads them in.
3. The example checks whether training is being continued from a checkpoint, and if so, gets the global step of the checkpoint and continues training from the last saved global step.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/1987?src=pr&el=h1) Report
> Merging [#1987](https://codecov.io/gh/huggingface/transformers/pull/1987?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/0cb163865a4c761c226b151283309eedb2b1ca4d?src=pr&el=desc) will **not change** coverage.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/1987/graphs/tree.svg?width=650&token=9qOlN6Hb1c&height=150&src=pr)](https://codecov.io/gh/huggingface/transformers/pull/1987?src=pr&el=tree)

```diff
@@           Coverage Diff           @@
##           master    #1987   +/-   ##
=======================================
  Coverage   82.67%   82.67%           
=======================================
  Files         111      111           
  Lines       16162    16162           
=======================================
  Hits        13362    13362           
  Misses       2800     2800
```



------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/1987?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/1987?src=pr&el=footer). Last update [0cb1638...bea4947](https://codecov.io/gh/huggingface/transformers/pull/1987?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Thank you for your contribution! I just had a quick look an will give it the time it deserves later. Were you aware of [Pytorch's guidelines](https://pytorch.org/tutorials/beginner/saving_loading_models.html#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training) when it comes to saving models? I reckon it might simplify your solution. 

Did this pull request check the epochs AND steps?
For instance, my training process stopped in epoch 2 step 6423.
If I run again it will continue from epoch 2 step 6423?

I am no expert, but looking at PyTorch schedulers' code they do keep the current global step in the state. The schedulers should thus continue from the last step. The thing I am wondering about is how to continue from the same data sample.

Hi guys,
Just wanted to ask this, Wouldn't too frequent caching to disks slow down the training overall?
We can have a flag added if the user wants to save every epoch, like ```file_name_{epoch}.pt```.
Plus we can save optimizer etc on the same weights file as well. 
Plus allowing users to specify those file names as well should be considered.
Thanks.

Hi,

@rlouf, 

Saving the all the checkpoints (model, tokenizer, optimizer, and scheduler) in one file like the pytorch example does would break the `from_pretrained` method, but I could change it to save the optimizer and scheduler in one file instead of two. 

I could change these lines:

```
# Saving
torch.save(optimizer.state_dict(), os.path.join(output_dir, 'optimizer.pt'))
torch.save(scheduler.state_dict(), os.path.join(output_dir, 'scheduler.pt'))

# Loading
optimizer.load_state_dict(torch.load(os.path.join(args.model_name_or_path, 'optimizer.pt')))
scheduler.load_state_dict(torch.load(os.path.join(args.model_name_or_path, 'scheduler.pt')))
```

to something like

```
# Saving
torch.save({
            'optimizer_state_dict': optimizer.state_dict(),
            'scheduler_state_dict': scheduler.state_dict()
            }, os.path.join(output_dir, 'training_state.pt'))

# Loading
checkpoint = torch.load(os.path.join(output_dir, 'training_state.pt'))
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
```

@marcoaleixo Yes, the code will resume training from the last saved checkpoint. The code saves the model every `--save_steps` training steps and saves the checkpoint in the format `checkpoint-global_step`, so we know exactly which global step the last checkpoint was on. From the global step we can figure out how many epochs have been trained and how many batches should be skipped in the current epoch to continue training from where we left off. The code for this is in [these](https://github.com/bkkaggle/transformers/blob/saving-and-resuming/examples/run_lm_finetuning.py#L230) lines

@rlouf To make sure you also continue from the last data sample, you would probably have to set a constant random seed in the dataloader and then `continue` through all the epochs and batches until you get to the saved checkpoint, which would take longer, especially if your dataset is very large.

@AdityaSoni19031997 You can set how often checkpoints are saved using the `--save_steps` flag. Saving the all the checkpoints (model, tokenizer, optimizer, and scheduler) in one file like the pytorch example does would break the `from_pretrained` method. Letting the user choose the file names might make it harder to automatically find and load in the files, and would require new command line parameters.

Tell me when this is not WIP anymore and ready to review, ok?

Hi @thomwolf, yes this is ready to review

@bkkaggle have you compared the training curves of a single run with the training curve of two half-runs with saving/reloading in the middle?

Hi @thomwolf, I ran some quick experiments to compare the loss curves. 

The loss curves are almost, but not exactly identical - probably because the `RandomSampler` doesn't accept a random seed.

One more thing: The mean loss at the end of a continued training run won't be the same as a training run completed in one go because the mean loss gets reset to 0 when continuing training. Do you want to find some way around this, maybe by also saving the running loss at each checkpoint? or would doing this add too much complexity for no benefit?

wandb dashboard with the loss curves: https://app.wandb.ai/bkkaggle/saving-and-resuming?workspace=default

- `vague_valley_29` is the original 1 epoch training run
- `vital_bird_30` is the same training run, but cancelled after step 100
-  `fresh-feather-34` is the training run resumed from `vital_bird_30`'s step 100 checkpoint

Ok, yes full determinism can be a rabbit-hole.
I think it's good to merge for me.
Ok for you as well @LysandreJik?

Yes, looks good to me!

I just finished updating the other pytorch examples: `run_xnli.py`, `run_ner.py`, `run_squad.py`, and `run_glue`. I pushed them up to my [branch](https://github.com/bkkaggle/transformers/commits/saving-and-resuming) and If you want, I can open another pull request for them.