Skip to content

support for wrapped schedulefree optimizer when using deepspeed - #3266

Merged
muellerzr merged 2 commits into
huggingface:mainfrom
winglian:deepspeed-schedulefree
Dec 2, 2024
Merged

support for wrapped schedulefree optimizer when using deepspeed#3266
muellerzr merged 2 commits into
huggingface:mainfrom
winglian:deepspeed-schedulefree

Conversation

@winglian

@winglian winglian commented Nov 26, 2024 •
edited
Loading

Copy link
Copy Markdown
Collaborator

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

BenjaminBossan commented Nov 28, 2024

Copy link
Copy Markdown
Member

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

@HuggingFaceDocBuilderDev

HuggingFaceDocBuilderDev commented Nov 28, 2024

Copy link
Copy Markdown

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
BenjaminBossan approved these changes Dec 2, 2024
View reviewed changes

@BenjaminBossan BenjaminBossan left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for this fix, LGTM.

muellerzr
muellerzr approved these changes Dec 2, 2024
View reviewed changes

@muellerzr muellerzr left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice, thanks for the fix!

@muellerzr
muellerzr merged commit 4a100ee into huggingface:main Dec 2, 2024
@winglian
winglian deleted the deepspeed-schedulefree branch December 2, 2024 22:21
@vineethsaivs vineethsaivs mentioned this pull request Jul 20, 2026
Merged
@XIONGPEILIN XIONGPEILIN mentioned this pull request Aug 11, 2026
Closed
2 of 4 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@BenjaminBossan BenjaminBossan BenjaminBossan approved these changes
+1 more reviewer
@muellerzr muellerzr muellerzr approved these changes
Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@winglian @BenjaminBossan @HuggingFaceDocBuilderDev @muellerzr