Skip to content

Commit 65b3dc4

Browse files
authored
scheduled removal of DeepSpeedPlugin.cpu_offload* parameters (#9244)
1 parent b046bd0 commit 65b3dc4

File tree

5 files changed

+10
-33
lines changed

5 files changed

+10
-33
lines changed

CHANGELOG.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -251,6 +251,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
251251
- Removed deprecated property `ModelCheckpoint.period` in favor of `ModelCheckpoint.every_n_epochs` ([#9213](https://github.com/PyTorchLightning/pytorch-lightning/pull/9213))
252252

253253

254+
- Removed deprecated properties `DeepSpeedPlugin.cpu_offload*` in favor of `offload_optimizer`, `offload_parameters` and `pin_memory` ([#9244](https://github.com/PyTorchLightning/pytorch-lightning/pull/9244))
254255

255256
### Fixed
256257

docs/source/advanced/advanced_gpu.rst

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,9 @@ You can also modify the ZeRO-Offload parameters via the plugin as below.
306306
307307
model = MyModel()
308308
trainer = Trainer(
309-
gpus=4, plugins=DeepSpeedPlugin(cpu_offload=True, allgather_bucket_size=5e8, reduce_bucket_size=5e8), precision=16
309+
gpus=4,
310+
plugins=DeepSpeedPlugin(offload_optimizer=True, allgather_bucket_size=5e8, reduce_bucket_size=5e8),
311+
precision=16,
310312
)
311313
trainer.fit(model)
312314
@@ -581,7 +583,7 @@ This saves memory when training larger models, however requires using a checkpoi
581583
gpus=4,
582584
plugins=DeepSpeedPlugin(
583585
stage=3,
584-
cpu_offload=True, # Enable CPU Offloading
586+
offload_optimizer=True, # Enable CPU Offloading
585587
cpu_checkpointing=True, # (Optional) offload activations to CPU
586588
),
587589
precision=16,
@@ -659,7 +661,7 @@ In some cases you may want to define your own DeepSpeed Config, to access all pa
659661
},
660662
"zero_optimization": {
661663
"stage": 2, # Enable Stage 2 ZeRO (Optimizer/Gradient state partitioning)
662-
"cpu_offload": True, # Enable Offloading optimizer state/calculation to the host CPU
664+
"offload_optimizer": True, # Enable Offloading optimizer state/calculation to the host CPU
663665
"contiguous_gradients": True, # Reduce gradient fragmentation.
664666
"overlap_comm": True, # Overlap reduce/backward operation of gradients for speed.
665667
"allgather_bucket_size": 2e8, # Number of elements to all gather at once.

pytorch_lightning/plugins/training_type/deepspeed.py

Lines changed: 1 addition & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@
3737
from pytorch_lightning.utilities.exceptions import MisconfigurationException
3838
from pytorch_lightning.utilities.imports import _DEEPSPEED_AVAILABLE
3939
from pytorch_lightning.utilities.types import LRSchedulerTypeTuple
40-
from pytorch_lightning.utilities.warnings import _warn, LightningDeprecationWarning, rank_zero_warn, WarningCache
40+
from pytorch_lightning.utilities.warnings import rank_zero_warn, WarningCache
4141

4242
warning_cache = WarningCache()
4343

@@ -124,9 +124,6 @@ def __init__(
124124
contiguous_memory_optimization: bool = False,
125125
synchronize_checkpoint_boundary: bool = False,
126126
load_full_weights: bool = False,
127-
cpu_offload: bool = False,
128-
cpu_offload_params: bool = False,
129-
cpu_offload_use_pin_memory: bool = False,
130127
) -> None:
131128
"""
132129
Provides capabilities to run training using the DeepSpeed library,
@@ -263,17 +260,6 @@ def __init__(
263260
"To use the DeepSpeed plugin, you must have DeepSpeed installed. pip install deepspeed"
264261
)
265262

266-
if cpu_offload or cpu_offload_params or cpu_offload_use_pin_memory:
267-
_warn(
268-
"The usage of `cpu_offload`, `cpu_offload_params`, and `cpu_offload_use_pin_memory` "
269-
"is deprecated since v1.4 and will be removed in v1.5."
270-
" From now on use `offload_optimizer`, `offload_parameters` and `pin_memory`.",
271-
category=LightningDeprecationWarning,
272-
)
273-
offload_optimizer = cpu_offload
274-
offload_parameters = cpu_offload_params
275-
pin_memory = cpu_offload_use_pin_memory
276-
277263
super().__init__(
278264
parallel_devices=parallel_devices,
279265
num_nodes=num_nodes,

tests/deprecated_api/test_remove_1-5.py

Lines changed: 0 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -16,10 +16,8 @@
1616

1717
from pytorch_lightning import Trainer
1818
from pytorch_lightning.core.decorators import auto_move_data
19-
from pytorch_lightning.plugins import DeepSpeedPlugin
2019
from tests.deprecated_api import no_deprecated_call
2120
from tests.helpers import BoringDataModule, BoringModel
22-
from tests.helpers.runif import RunIf
2321

2422

2523
def test_v1_5_0_auto_move_data():
@@ -43,16 +41,6 @@ def test_v1_5_0_datamodule_setter():
4341
assert any("The `LightningModule.datamodule`" in w for w in warning_cache)
4442

4543

46-
@RunIf(deepspeed=True)
47-
@pytest.mark.parametrize(
48-
"params", [dict(cpu_offload=True), dict(cpu_offload_params=True), dict(cpu_offload_use_pin_memory=True)]
49-
)
50-
def test_v1_5_0_deepspeed_cpu_offload(tmpdir, params):
51-
52-
with pytest.deprecated_call(match="is deprecated since v1.4 and will be removed in v1.5"):
53-
DeepSpeedPlugin(**params)
54-
55-
5644
def test_v1_5_0_distributed_backend_trainer_flag():
5745
with pytest.deprecated_call(match="has been deprecated and will be removed in v1.5."):
5846
Trainer(distributed_backend="ddp_cpu")

tests/plugins/test_deepspeed_plugin.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -372,13 +372,13 @@ def test_deepspeed_custom_activation_checkpointing_params(tmpdir):
372372

373373
@RunIf(min_gpus=1, deepspeed=True)
374374
def test_deepspeed_assert_config_zero_offload_disabled(tmpdir, deepspeed_zero_config):
375-
"""Ensure if we use a config and turn off cpu_offload, that this is set to False within the config."""
375+
"""Ensure if we use a config and turn off offload_optimizer, that this is set to False within the config."""
376376

377-
deepspeed_zero_config["zero_optimization"]["cpu_offload"] = False
377+
deepspeed_zero_config["zero_optimization"]["offload_optimizer"] = False
378378

379379
class TestCallback(Callback):
380380
def on_before_accelerator_backend_setup(self, trainer, pl_module) -> None:
381-
assert trainer.training_type_plugin.config["zero_optimization"]["cpu_offload"] is False
381+
assert trainer.training_type_plugin.config["zero_optimization"]["offload_optimizer"] is False
382382
raise SystemExit()
383383

384384
model = BoringModel()

0 commit comments

Comments
 (0)