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.

Sign up for GitHub

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
merged 3 commits into from
Jul 8, 2024
Merged

Update DeepSpeedOptimizer import for deepspeed >= 0.14.1 #20040

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
401b4fa
[FIX] fix deepspeed DeepSpeedOptimizer import error
linbo0518 Jul 3, 2024
8b5a5ec
[FIX] fix deepspeed DeepSpeedOptimizer import error when deepspeed ve…
linbo0518 Jul 8, 2024
6071957
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Jul 8, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion src/lightning/fabric/strategies/deepspeed.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@
from deepspeed import DeepSpeedEngine

_DEEPSPEED_AVAILABLE = RequirementCache("deepspeed")
_DEEPSPEED_GREATER_EQUAL_0_14_1 = RequirementCache("deepspeed>=0.14.1")


# TODO(fabric): Links in the docstrings to PL-specific deepspeed user docs need to be replaced.
Expand Down Expand Up @@ -498,7 +499,10 @@ def load_checkpoint(
)
engine = engines[0]

from deepspeed.runtime import DeepSpeedOptimizer
if _DEEPSPEED_GREATER_EQUAL_0_14_1:
from deepspeed.runtime.base_optimizer import DeepSpeedOptimizer
else:
from deepspeed.runtime import DeepSpeedOptimizer

optimzer_state_requested = any(isinstance(item, (Optimizer, DeepSpeedOptimizer)) for item in state.values())

Expand Down
Loading