Introduce Stateful Callbacks
Author: muellerzr
Pull request: yes

# What does this PR do?

This PR builds a foundation for stateful callbacks inside `Trainer`. Right now these are isolated to the core callbacks that exist, but can be expanded upon as needed later.

Specifically a user must enable `restore_callback_states` in the `TrainingArguments` to enable this behavior.

## System Design:

To keep with not having to deal with pickles/interact with the `TrainerState` properly (where callback data should be stored, likely), we need a way to recreate the exact state of a `Callback`. To do so callbacks should implement a `save_state` function which will create a `dict` of `args` and `attributes` that should be set.

From this we can then recreate the constructor of an existing callback with these new options. (this is done in case there is logic in a callback's `__init__` that could be needed).

For example, here is the one for `EarlyStoppingCallback`:

```python
        def save(self) -> dict:
            return {
                "args": {
                    "early_stopping_patience": self.early_stopping_patience,
                    "early_stopping_threshold": self.early_stopping_threshold,
                },
                "attributes": {
                    "early_stopping_patience_counter": self.early_stopping_patience_counter,
                }
            }
```

## What about Callbacks that I forget to add back in?

As this relies on a similar **state**, I chose to keep aligned with how we deal with states/checkpointing in Accelerate, wherein users must recreate the exact same initial scenario when resuming training.

Or, in other words:

If you included the `EarlyStoppingCheckpoint` initially, but when you do `resume_from_checkpoint` you did not include that callback, we do not magically init and resume those states. Instead we will simply keep going and give you a nice warning message saying that things went amiss/some callback states weren't resumed.

## Limitations

As we're maintaining that these exist in the `TrainerState` (which again, makes sense for us to use this here), items should be JSON-serializable. 

The aim here is most likely not many callbacks actually *need* this, so only a few as needed over time can be added to this without much complexity.


Fixes https://github.com/huggingface/transformers/issues/28544


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@amyeroberts 

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29666). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@amyeroberts let me know if we need more docs or if I should rewrite the code a bit if it feels to complex/magical :)

@amyeroberts I'd like one more review please, just to make sure the multiple-versions logic makes sense to you!

Hi,  it looks like the state of stateful callbacks are not updated before saving state to trainer_state.json on  _save_checkpoint method:

https://github.com/huggingface/transformers/blob/fe008d6ebea1f5770b740991daeefd9322fa434a/src/transformers/trainer.py#L2917-L2920
Only TrainerControl state is saved, instead of getting all ExportableState callbacks like here:
https://github.com/huggingface/transformers/blob/fe008d6ebea1f5770b740991daeefd9322fa434a/src/transformers/trainer.py#L674-L675

So it should be:
```
if self.args.should_save: 
    for cb in [cb for cb in self.callback_handler.callbacks  + [self.control] if isinstance(cb, ExportableState)]:
        self.state.stateful_callbacks[cb.__class__.__name__] = cb.state()
    self.state.save_to_json(os.path.join(output_dir, TRAINER_STATE_NAME))    
```
Do you agree?