Add support for GrokAdamW optimizer
Author: ehartford
Pull request: yes

# What does this PR do?

### Add support for GrokAdamW optimizer

This PR adds support for the GrokAdamW optimizer to the `transformers` library.

#### Changes Introduced
- Integrated the GrokAdamW optimizer into the `Trainer` class.
- Added error handling to prompt users to install the `grokadamw` package if not already installed.

#### Motivation
The GrokAdamW optimizer enhances training performance and stability for certain models, providing users with more optimization options.

#### Dependencies
- **grokadamw**: Users need to install this package via `pip install grokadamw`.

#### Code Changes
- **trainer.py**: Added a new conditional block to import and configure the GrokAdamW optimizer.

```python
elif args.optim == "grokadamw":
    try:
        from grokadamw import GrokAdamW

        optimizer_cls = GrokAdamW
        optimizer_kwargs.update(
            {
                "alpha_init": float(optim_args.get("alpha_init", 0.98)),
                "lamb": float(optim_args.get("lamb", 2.0)),
                "gamma": float(optim_args.get("gamma", 0.1)),
                "grokking_signal_decay_rate": float(optim_args.get("grokking_signal_decay_rate", 0.1)),
                "gradient_clipping": float(optim_args.get("gradient_clipping", 1.0)),
            }
        )
    except ImportError:
        raise ValueError("Please install grokadamw with `pip install grokadamw`")
```

#### Testing
- Verified the integration with a test script to ensure the optimizer works as expected.

## 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).
- [ ] Did you write any new necessary tests?


## Who can review?

@muellerzr and @SunMarc


Comments:
@amyeroberts I do not, best we wait until @LysandreJik is back for that question! (As long as it's not too problematic, I don't see an issue with adding more optimizers)

> Thanks for adding this! Is there an associated issue to add this as a feature request?
> 
> I'll let @muellerzr and @SunMarc chip in on whether this is something we want to add. One thing to note is that this implementation is under the MIT license. @muellerzr do you know how we typically handle difference licensing for integrations like this?
> 
> As commented - tests would need to be added - the LOMO PR is a good reference here

There is not an associated issue.  I am hoping to add it, in order to make it more accessible to users.

I have updated the license to Apache 2.0 to accommodate your license.

I added the tests as requested.


What do you guys think about the grokking functions?  As it is, I don't see how they can easily pass those in.  And so I just use the default grokking function (which honestly 99.9% of people will do). Otherwise they would need to pass the function as a string which would then be eval'd - fugly.

But this seems a good compromise for simplicity sake and if they want custom grokking functions they can use it manually.

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

@ehartford currently they'd need to pass them in via `optim_args`, but since many optimizers are getting complex and needing special `TrainingArgument` values, I'd be open to a follow-up PR which refactors things to include `optim_args` and `optim_kwargs` as part of `TrainingArguments` which then get passed to the optimizer init

Hello,
May I ask please, did you need anything else from me? 
Is there another review required?

@ehartford running this i get hit with this error. can you check 

Thanks for the contribution anyways

```
  File "/home/shas/qbFinetuning/processor/LLM/finetunerbasemodel.py", line 69, in run_full_finetuningjob
    self._train()
  File "/home/shas/qbFinetuning/processor/LLM/error_utils.py", line 82, in wrapper_exit
    return func(*args, **kwargs)
  File "/home/shas/qbFinetuning/processor/LLM/error_utils.py", line 55, in wrapper_retry
    return func(*args, **kwargs)
  File "/home/shas/qbFinetuning/processor/LLM/finetunerbasemodel.py", line 54, in _train
    return self.train(*args,**kwargs)
  File "/home/shas/.local/lib/python3.10/site-packages/trl/trainer/sft_trainer.py", line 451, in train
    output = super().train(*args, **kwargs)
  File "/home/shas/miniforge3/lib/python3.10/site-packages/transformers/trainer.py", line 1963, in train
    return inner_training_loop(
  File "/home/shas/.local/lib/python3.10/site-packages/accelerate/utils/memory.py", line 146, in decorator
    return function(batch_size, *args, **kwargs)
  File "/home/shas/miniforge3/lib/python3.10/site-packages/transformers/trainer.py", line 2366, in _inner_training_loop
    self.optimizer.step()
  File "/home/shas/.local/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 130, in wrapper
    return func.__get__(opt, opt.__class__)(*args, **kwargs)
  File "/home/shas/.local/lib/python3.10/site-packages/accelerate/optimizer.py", line 170, in step
    self.optimizer.step(closure)
  File "/home/shas/.local/lib/python3.10/site-packages/torch/optim/lr_scheduler.py", line 130, in wrapper
    return func.__get__(opt, opt.__class__)(*args, **kwargs)
  File "/home/shas/.local/lib/python3.10/site-packages/torch/optim/optimizer.py", line 484, in wrapper
    out = func(*args, **kwargs)
  File "/home/shas/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
  File "/home/shas/.local/lib/python3.10/site-packages/grokadamw/grokadamw.py", line 53, in step
    return self._step_impl(closure, use_amp=False)
  File "/home/shas/.local/lib/python3.10/site-packages/grokadamw/grokadamw.py", line 92, in _step_impl
    _apply_updates()
  File "/home/shas/.local/lib/python3.10/site-packages/grokadamw/grokadamw.py", line 86, in _apply_updates
    self._update_group(group, params_with_grad, grads, grokking_signal)
  File "/home/shas/.local/lib/python3.10/site-packages/grokadamw/grokadamw.py", line 132, in _update_group
    state = group['state'][p]
KeyError: 'state'
  0%|                     

```

Please update GrokAdamW 
I have version 0.1.2 

Thanks a ton , working now!