pytorch-lightning
Update DeepSpeedOptimizer import for deepspeed >= 0.14.1
#20040
Merged
Overview
Commits
3
Changes
View On GitHub

Update DeepSpeedOptimizer import for deepspeed >= 0.14.1 #20040

awaelchli merged 3 commits into Lightning-AI:master from linbo0518:fix/deepspeed_optimizer_import_error
linbo0518
linbo0518320 days ago (edited 320 days ago)

What does this PR do?

Fix DeepSpeedOptimizer import error

DeepSpeed team remove DeepSpeedOptimizer from the __init__.py

deepspeedai/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/

linbo0518 [FIX] fix deepspeed DeepSpeedOptimizer import error
401b4fa5
linbo0518 linbo0518 requested a review from lantiga lantiga 320 days ago
linbo0518 linbo0518 requested a review from Borda Borda 320 days ago
linbo0518 linbo0518 requested a review from tchaton tchaton 320 days ago
linbo0518 linbo0518 requested a review from awaelchli awaelchli 320 days ago
linbo0518 linbo0518 requested a review from justusschock justusschock 320 days ago
github-actions github-actions added fabric
awaelchli
awaelchli commented on 2024-07-06
awaelchli317 days ago

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

Conversation is marked as resolved
Show resolved
src/lightning/fabric/strategies/deepspeed.py
499499 engine = engines[0]
500500
501 from deepspeed.runtime import DeepSpeedOptimizer
501
 from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
awaelchli317 days ago

Since this was a recent breaking change in deepspeed, if I see correctly in version 0.14.1, let's make the import conditioned on the version being used:

In L46, define:

_DEEPSPEED_GREATER_EQUAL_0_14_1 = RequirementCache("deepspeed>=0.14.1")

Then in the import here:

Suggested change
from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
if _DEEPSPEED_GREATER_EQUAL_0_14_1:
from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
else:
from deepspeed.runtime import DeepSpeedOptimizer
linbo0518315 days ago

Great

awaelchli awaelchli changed the title [FIX] fix deepspeed DeepSpeedOptimizer import error Update DeepSpeedOptimizer import for deepspeed >= 0.14.1 317 days ago
awaelchli awaelchli added strategy: deepspeed
awaelchli awaelchli added community
awaelchli awaelchli added this to the 2.4 milestone 317 days ago
awaelchli awaelchli removed this from to the 2.4 milestone 317 days ago
awaelchli awaelchli added this to the 2.3.x milestone 317 days ago
codecov
codecov317 days ago (edited 225 days ago)

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 73 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
linbo0518 [FIX] fix deepspeed DeepSpeedOptimizer import error when deepspeed ve…
8b5a5ec7
pre-commit-ci[bot] [pre-commit.ci] auto fixes from pre-commit.com hooks
6071957a
awaelchli
awaelchli approved these changes on 2024-07-08
awaelchli awaelchli merged 9e6c02a9 into master 315 days ago
awaelchli
awaelchli315 days ago😄 1

Thank you for doing this @linbo0518 🚀

Login to write a write a comment.

Login via GitHub

Reviewers
Assignees
No one assigned
Labels
strategy: deepspeed fabric community
Milestone
2.3.x
Terms Of Use Privacy Policy Refund Policy Impressum