Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update DeepSpeedOptimizer import for deepspeed >= 0.14.1 #20040

Merged

Conversation

linbo0518
Copy link
Contributor

@linbo0518 linbo0518 commented Jul 3, 2024

What does this PR do?

Fix DeepSpeedOptimizer import error

DeepSpeed team remove DeepSpeedOptimizer from the __init__.py

microsoft/DeepSpeed@c56a4b9

[rank1]: Traceback (most recent call last):
[rank1]:   File "/data/projects/livdet2025/train.py", line 140, in <module>
[rank1]:     main()
[rank1]:   File "/data/projects/livdet2025/train.py", line 134, in main
[rank1]:     train_helper.load_state(checkpoint, model, optimizer, lr_scheduler)
[rank1]:   File "/data/projects/livdet2025/utils.py", line 253, in load_state
[rank1]:     }
[rank1]:   File "/data/.pyenv/versions/livdet/lib/python3.10/site-packages/lightning/fabric/fabric.py", line 775, in load
[rank1]:     remainder = self._strategy.load_checkpoint(path=path, state=unwrapped_state, strict=strict)
[rank1]:   File "/data/.pyenv/versions/livdet/lib/python3.10/site-packages/lightning/fabric/strategies/deepspeed.py", line 501, in load_checkpoint
[rank1]:     from deepspeed.runtime import DeepSpeedOptimizer
[rank1]: ImportError: cannot import name 'DeepSpeedOptimizer' from 'deepspeed.runtime' (/data/.pyenv/versions/livdet/lib/python3.10/site-packages/deepspeed/runtime/__init__.py)

📚 Documentation preview 📚: https://pytorch-lightning--20040.org.readthedocs.build/en/20040/

@github-actions github-actions bot added the fabric lightning.fabric.Fabric label Jul 3, 2024
Copy link
Member

@awaelchli awaelchli left a comment

Choose a reason for hiding this comment

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

Looks good @linbo0518, thanks! We still have to make it compatible with other versions, see my suggestion below.

src/lightning/fabric/strategies/deepspeed.py Outdated Show resolved Hide resolved
@awaelchli awaelchli changed the title [FIX] fix deepspeed DeepSpeedOptimizer import error Update DeepSpeedOptimizer import for deepspeed >= 0.14.1 Jul 6, 2024
@awaelchli awaelchli added strategy: deepspeed community This PR is from the community labels Jul 6, 2024
@awaelchli awaelchli modified the milestones: 2.4, 2.3.x Jul 6, 2024
Copy link

codecov bot commented Jul 6, 2024

Codecov Report

Attention: Patch coverage is 75.00000% with 1 line in your changes missing coverage. Please review.

Project coverage is 90%. Comparing base (693c21a) to head (6071957).
Report is 7 commits behind head on master.

Additional details and impacted files
@@            Coverage Diff             @@
##           master   #20040      +/-   ##
==========================================
+ Coverage      83%      90%      +7%     
==========================================
  Files         427      266     -161     
  Lines       35460    22942   -12518     
==========================================
- Hits        29596    20722    -8874     
+ Misses       5864     2220    -3644     

@awaelchli awaelchli merged commit 9e6c02a into Lightning-AI:master Jul 8, 2024
94 of 96 checks passed
@awaelchli
Copy link
Member

Thank you for doing this @linbo0518 🚀

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
community This PR is from the community fabric lightning.fabric.Fabric strategy: deepspeed
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants