support for wrapped schedulefree optimizer when using deepspeed

[winglian]

What does this PR do?

Axolotl community reported an issue with schedule free AdamW with deepspeed:

[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/transformers/trainer.py", line 2126, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/transformers/trainer.py", line 2484, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/transformers/trainer.py", line 3617, in training_step
[rank0]:     self.accelerator.backward(loss, **kwargs)
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/accelerate/accelerator.py", line 2233, in backward
[rank0]:     self.deepspeed_engine_wrapped.backward(loss, **kwargs)
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/accelerate/utils/deepspeed.py", line 195, in backward
[rank0]:     self.engine.step()
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2213, in step
[rank0]:     self._take_model_step(lr_kwargs)
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/deepspeed/runtime/engine.py", line 2119, in _take_model_step
[rank0]:     self.optimizer.step()
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1905, in step
[rank0]:     self._optimizer_step(i)
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 1811, in _optimizer_step
[rank0]:     self.optimizer.step()
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/torch/optim/lr_scheduler.py", line 75, in wrapper
[rank0]:     return wrapped(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/torch/optim/optimizer.py", line 391, in wrapper
[rank0]:     out = func(*args, **kwargs)
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/root/miniconda3/envs/py3.11/lib/python3.11/site-packages/schedulefree/adamw_schedulefree.py", line 112, in step
[rank0]:     raise Exception("Optimizer was not in train mode when step is called. "
[rank0]: Exception: Optimizer was not in train mode when step is called. Please insert .train() and .eval() calls on the optimizer. See documentation for details.

Fixes # (issue)

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline,
  Pull Request section?
• Was this discussed/approved via a Github issue or the forum? 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, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

@muellerzr

[BenjaminBossan]

Thanks @winglian. I think it would be helpful to add a comment on why checking for optimizer.optimizer is needed.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[BenjaminBossan]

Thanks for this fix, LGTM.

[muellerzr]

Nice, thanks for the fix!