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

Support shortcut name for DeepSpeed stage 1 offload #19075

7 changes: 7 additions & 0 deletions src/lightning/fabric/strategies/deepspeed.py
Original file line number Diff line number Diff line change
Expand Up @@ -551,6 +551,13 @@ def clip_gradients_value(
def register_strategies(cls, strategy_registry: _StrategyRegistry) -> None:
strategy_registry.register("deepspeed", cls, description="Default DeepSpeed Strategy")
strategy_registry.register("deepspeed_stage_1", cls, description="DeepSpeed with ZeRO Stage 1 enabled", stage=1)
strategy_registry.register(
"deepspeed_stage_1_offload",
cls,
description="DeepSpeed with ZeRO Stage 1 and optimizer CPU Offload",
stage=1,
offload_optimizer=True,
)
strategy_registry.register("deepspeed_stage_2", cls, description="DeepSpeed with ZeRO Stage 2 enabled", stage=2)
strategy_registry.register(
"deepspeed_stage_2_offload",
Expand Down
1 change: 1 addition & 0 deletions src/lightning/pytorch/CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
- Added a utility function and CLI to consolidate FSDP sharded checkpoints into a single file ([#19213](https://github.com/Lightning-AI/lightning/pull/19213))
- The TQDM progress bar now respects the env variable `TQDM_MINITERS` for setting the refresh rate ([#19381](https://github.com/Lightning-AI/lightning/pull/19381))
- Added support for saving and loading stateful training DataLoaders ([#19361](https://github.com/Lightning-AI/lightning/pull/19361))
- Added shortcut name `strategy='deepspeed_stage_1_offload'` to the strategy registry ([#19075](https://github.com/Lightning-AI/lightning/pull/19075))

### Changed

Expand Down
1 change: 1 addition & 0 deletions tests/tests_fabric/strategies/test_registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ def test_available_strategies_in_registry():
"ddp",
"deepspeed",
"deepspeed_stage_1",
"deepspeed_stage_1_offload",
"deepspeed_stage_2",
"deepspeed_stage_2_offload",
"deepspeed_stage_3",
Expand Down
Loading