Add AdEMAMix optimizer
Author: matthewdouglas
Pull request: yes

# What does this PR do?

Adds support in `Trainer` and `TrainingArguments` for the bitsandbytes implementations of the AdEMAMix optimizer as proposed here: https://hf.co/papers/2409.03137

Compared to AdamW, there is an additional "slow" momentum term to keep older gradients relevant. This authors show that this leads to faster convergence and reduces forgetting during training.

Support for this new optimizer was introduced in [bitsandbytes v0.44.0](https://github.com/bitsandbytes-foundation/bitsandbytes/releases/tag/v0.44.0). This implementation balances out the additional memory overhead when compared to Adam by providing 8bit and 32bit versions that support paging.

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->

Fixes #33392


## 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?

@SunMarc 

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

Merging as failing hub tests are unrelated to this PR 

@matthewdouglas @amyeroberts Thanks for creating this, just a suggestion, however, it would be worth setting 'T<sub>alpha</sub> = T<sub>β3</sub> = T = number of training steps' by default, as recommended by the AdEMAMIx paper, page 6:
<img width="790" alt="image" src="https://github.com/user-attachments/assets/41d2baac-95b6-46dd-8e90-0a27d62ca4f4">

I tried using this code to train a model not realising that those hyperparameters were not being set by default by `bitsandbytes`, which may explain why I was expericing instability: page 29.

<img width="819" alt="image" src="https://github.com/user-attachments/assets/6d6214af-d809-4980-ace8-af1d82cc7629">

Good defaults has always been what sets Hugging Face apart 😁 

Hi @umarbutler, thanks for flagging! That seems reasonable, especially as we already set defaults for e.g. beta3. Would you like to open a PR to add this? 

@amyeroberts I believe this issue on default values for `t_alpha` and `t_beta3` is still not resolved. Can you confirm? As I cant find anything on default values being set in the repository. I can raise a PR to fix this.