Deprecates AdamW and adds `--optim`
Author: manuelciosici
Pull request: yes

# What does this PR do?

<!--
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 #14539


## 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/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] 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/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?


## Who can review?

@stas00 

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- albert, bert, xlm: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @n1t0, @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

Documentation: @sgugger

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Misc
@stas00 
 * I added `FusedAdam` based on your comment in #14539
 * Since both `--optim adafactor` and `--adafactor` server the same purpose, I marked `--adafactor` as deprecated. I copy-pasted a deprecation warning that mentions `transformers` version 5. Let me know if the deprecation warning should say something else.
 * Let me know if I missed anything else.

Comments:
@sgugger Thank you for the suggestions. I think an enum makes sense. I will add that and fix the missing `optimizer_cls` line.

@manuelciosici. have we lost you or are you just on an extended vacation? 

Let's finish this up, so we can test out various optimizers. I will sync the recent changes in doc formatting.

I was just told that `apex.optimizers.FusedAdam` is an even faster fused optimizer than torch's - we added it already but going to benchmark it.

update: see updated benchmarks here:

1. [RTX-3090](https://github.com/huggingface/transformers/issues/14608#issuecomment-1005219385)
2. [A100](https://github.com/huggingface/transformers/issues/15026#issuecomment-1005220263)

------------

I'm working on a [neat HF Trainer benchmarking tool](https://github.com/huggingface/transformers/pull/14934), so here it is applied to the changes introduced by this PR:

| Variation               |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:------------------------|------------------------------------:|------------:|----------------:|
| --optim adamw_hf        |                             117.544 |         32 |         2.19851 |
| --optim adamw_torch     |                             112.951 |         27 |         2.19829 |
| --optim adafactor       |                              89.194 |         0 |         2.20484 |
| --optim apex_fused_adam |                             126.232 |         42 |         2.19832 |

So torch's AdamW appears to be even slower than ours. So clearly apex's AdamW is the way to go speed-wise.

Note, that the absolute and relative results will be different on a different GPU and a different finetuning setup, but most likely the current fastest optimizer will remain fastest, etc. 

Reproducibility and other info:

```
Datetime    : 2021-12-28 20:56:42

Software:
transformers: 4.16.0.dev0
torch       : 1.10.1
cuda        : 11.3
python      : 3.8.11

Hardware:
1 GPUs      : NVIDIA GeForce RTX 3090, 23.70GB

The benchmark command line was:

CUDA_VISIBLE_DEVICES=0 python \
/hf/transformers-trainer-benchmark/scripts/benchmark/trainer-benchmark.py \
--base-cmd \
' \
examples/pytorch/translation/run_translation.py --model_name_or_path t5-base --output_dir output_dir \
--do_train --label_smoothing 0.1 --logging_strategy no --save_strategy no --per_device_train_batch_size 16 \
--max_source_length 512 --max_target_length 512 --num_train_epochs 1 --overwrite_output_dir \
--source_lang en --target_lang ro --dataset_name wmt16 --dataset_config "ro-en" \
--source_prefix "translate English to Romanian: "  --warmup_steps 50 \
--max_train_samples 10000 --dataloader_num_workers 2 \
' \
--target-metric-key train_samples_per_second --repeat-times 1 --variations \
'--optim adamw_hf|--optim adamw_torch|--optim adafactor|--optim apex_fused_adam' \
--table-format github --report-metric-keys train_loss

```

@stas00 You have not lost me. I've been on vacation and before that, I accidentally messed up my GitHub notification settings, so I didn't know you had also reviewed the code. I applied all the suggestions from you.

I also changed to an `Enum` for optimizer string values as @sgugger suggested.

Let me know if I should change anything else.

@stas00 Let me know if I should squash all the commits to a single commit so they don't pollute `master`'s commit log.

> @stas00 You have not lost me. [...]

Great! I will review your edits today and we can finalize it then.

> @stas00 Let me know if I should squash all the commits to a single commit so they don't pollute `master`'s commit log.

No need to. We always squash on merge

I implemented @stas00 's suggestions (except for the `warnings.warn` bit, see the discussion above). I also remembered that @sgugger suggested adding a unit test for the optimizer selection, so I added that too.

@stas00 I managed to parameterize the tests for AdamW (HF and torch) and Adafactor and refactored the testing code.

For testing `apex.optimizers.FusedAdam`, since `Trainer.get_optimizer_cls_and_kwargs` does not use `FusedAdam` but only returns a class called `apex.optimizers.FusedAdam`, I thought I could mock `FusedAdam`'s  presence or absence independent of an `apex` installation.

In `test_fused_adam`, I mock `apex.optimizers.FusedAdam` and check that `Trainer.get_optimizer_cls_and_kwargs` returns the object mocking `apex.optimizers.FusedAdam`.

In `test_fused_adam_no_apex`, I simulate a missing `apex.optimizers.FusedAdam` by setting the `apex` namespace to `None` so `Trainer.get_optimizer_cls_and_kwargs` cannot import `apex.optimizers.FusedAdam` even if `apex` is installed.

Does this testing approach make sense? Or should I write `apex`-dependent tests and use `is_apex_available`?


> @stas00 I managed to parameterize the tests for AdamW (HF and torch) and Adafactor and refactored the testing code.

Awesome! much better!

> For testing `apex.optimizers.FusedAdam`, since `Trainer.get_optimizer_cls_and_kwargs` does not use `FusedAdam` but only returns a class called `apex.optimizers.FusedAdam`, I thought I could mock `FusedAdam`'s presence or absence independent of an `apex` installation.

Perhaps we are looking at different code?  `Trainer.get_optimizer_cls_and_kwargs` branches are almost identical for torch or apex. What do you mean it only returns a class. It returns kwargs too.

> In `test_fused_adam`, I mock `apex.optimizers.FusedAdam` and check that `Trainer.get_optimizer_cls_and_kwargs` returns the object mocking `apex.optimizers.FusedAdam`.
> 
> In `test_fused_adam_no_apex`, I simulate a missing `apex.optimizers.FusedAdam` by setting the `apex` namespace to `None` so `Trainer.get_optimizer_cls_and_kwargs` cannot import `apex.optimizers.FusedAdam` even if `apex` is installed.
> 
> Does this testing approach make sense? Or should I write `apex`-dependent tests and use `is_apex_available`?

Your testing is fabulous, Manuel. But you're not testing the real thing when it's available. What if apex does a breaking change and we are still mocking what it should be doing and not what it does. Hence, I'm suggesting that there should be an identical test to torch's except need to inject a skip in that parameterized test if  `testing apex and not is_apex_available()`. Would it'd be easier if I wrote it?


@stas00 

>  What do you mean it only returns a class. It returns kwargs too.

You are right, it returns a class and the keyword arguments. I didn't express myself properly.

I added a new test that uses `FusedAdam` from `apex` if `is_apex_available` is `True`. I couldn't figure out how to test using a parameterization of `test_supported_optim` since I couldn't figure out how to add the `FusedAdam` class to `test_params` without causing the code to throw a `ModuleNotFoundError` in environments without `apex`. In other words, I can't figure out how to conditionally add to `test_params` so that `@parameterized.expand` still works. If there is a way, can you show me?

Done.

Also added an actual `train()` in the test, which is the ultimate test.  We typically don't test API at such low level when it gets functionally tested via a higher end API. The functional `train()` test does it all by the fact that it should succeed and all the testing is already done there automatically.

Note that I get the last test fail:

```
tests/test_trainer.py .F

___________________________________________________ TrainerOptimizerChoiceTest.test_fused_adam_no_apex ___________________________________________________

self = <tests.test_trainer.TrainerOptimizerChoiceTest testMethod=test_fused_adam_no_apex>

    def test_fused_adam_no_apex(self):
        args = TrainingArguments(optim=OptimizerNames.ADAM_APEX_FUSED.value, output_dir="None")
    
        # Pretend that apex does not exist, even if installed. By setting apex to None, importing
        # apex will fail even if apex is installed.
        with patch.dict("sys.modules", {"apex": None}):
            with self.assertRaises(ValueError):
>               Trainer.get_optimizer_cls_and_kwargs(args)
E               AssertionError: ValueError not raised

tests/test_trainer.py:1795: AssertionError

tests/test_trainer.py .....   
```

I'm not an expert on Mock, so not sure why it doesn't override like it should.

I have apex installed when I run this test.

Can you reproduce it?

update: fixed it. Needed `with patch.dict("sys.modules", {"apex.optimizers": None}):` since it's importing from `apex.optimizers` and not apex.

@stas00 Thank you for the code cleanup, figuring out the mocking issue, and for your patience. This PR has been an educational experience (I didn't know about `parameterized`). I'm looking forward to figuring out what my next contribution should be. Let me know if you have any suggestions.

Since clearly you're interesting in easier testing, you may find some useful tidbits in this extensive doc: https://huggingface.co/docs/transformers/testing, e.g. for different parameterization situations https://huggingface.co/docs/transformers/testing#parametrization

optimizers-wise I think the next interesting but challenging thing to add is BNB 8bit optimizer
https://github.com/huggingface/transformers/issues/14819 but we are still discussing how it'd work.

The other thing to potentially experiment with is 
https://www.deepspeed.ai/tutorials/onebit-adam/
but I haven't had a chance to understand it so I have no idea whether it can be used outside of Deepspeed or just with deepspeed.

I think the only other remaining item that I didn't hear you weigh on, @sgugger, is whether we should call these `--adam_foo` or `--adamw_foo` since the class names are `AdamW` (except apex).

Oh sorry, I didn't catch that. It should adamw everywhere IMO.

> Oh sorry, I didn't catch that. It should adamw everywhere IMO.

Thank you for validating that, @sgugger. 

@manuelciosici, so one more tiny change please `s|--adam_|--adamw_|` 

Thank you!


I removed `.value` from everywhere and ensured that the tests pass. I have also changed optimizer name strings as @stas00  asked.

Let me know if I should change anything else.

but the rest needs to updated to match, e.g. currently many torch tests fail with:
```
E           ValueError: Trainer cannot instantiate unsupported optimizer: adamw_hf
```

@stas00 I just saw that. I'm trying to figure out what I misunderstood.

@stas00 I surprised that fixes it. On my end, I just fixed it by adding `self.optim = OptimizerNames(self.optim)` in post-init. I also had to remove a now redundant unit test.

It's the same as:
```
elif args.optim == OptimizerNames.ADAMW_HF:
```

I'm not sure if you can see the failing tests from CI, so I thought it'd be faster to push in the fix as I saw where it was failing.

Are you still planning to push a change? You said removing a unit test.

Let us know when you're done.


@stas00 Commit e73249ce130bc81c392f5b9a3224f5ba64a2214d makes the unit tests pass, but it doesn't work when `--optim` is explicitly set on the command line. `TrainingArguments` does not automatically convert strings to enums. For a parallel, see https://github.com/huggingface/transformers/blob/9a94bb8e218033cffa1ef380010b528410ba3ca7/src/transformers/training_args.py#L755

So, calling a script on the command line (for example, `examples/pytorch/language-modeling/run_clm.py`) with any explicit `--optim` (even if only with `--optim adamw_hf`), throws an error from `Trainer.get_optimizer_cls_and_kwargs` since `get_optimizer_cls_and_kwargs` receives a string that does not match any of the if branches which test for `Enum` values.

I added another test for specifying the optimizer name as a string. Also, I removed `test_optim_unsupported` since, with these commits, `--optim` no longer accepts strings that are not in `OptimizerNames`.

I also changed the default value from `OptimizerNames.ADAMW_HF` to `OptimizerNames.ADAMW_HF.value`. With `default=OptimizerNames.ADAMW_HF` calling `--help` on the CLI gives:

```
  --optim {adamw_hf,adamw_torch,adamw_apex_fused,adafactor}
                        The optimizer to use. (default:
                        OptimizerNames.ADAMW_HF)
```

While `default=OptimizerNames.ADAMW_HF.value` gives

```
  --optim {adamw_hf,adamw_torch,adamw_apex_fused,adafactor}
                        The optimizer to use. (default: adamw_hf)
```

The first one leaks the internal object name, while the second indicates the string we want users to pass.

Finally, I make `OptimizerNames` inherit `ExplicitEnum` instead of `Enum` because I saw `SchedulerType` do the same and it seems more elegant.
https://github.com/huggingface/transformers/blob/9a94bb8e218033cffa1ef380010b528410ba3ca7/src/transformers/trainer_utils.py#L281


I second that - thank you, @manuelciosici!

@stas00 @sgugger Thank you for guiding me!