From 2a4b251c6290b8e71a04df98e5e6f49837077356 Mon Sep 17 00:00:00 2001 From: Masahiro Wada Date: Sun, 3 Jul 2022 05:43:01 +0000 Subject: [PATCH 1/9] Add finetuning strategies for DeepSpeed --- flash/core/finetuning.py | 52 +++++++++++++++++++++++++++++- flash/core/utilities/imports.py | 1 + tests/core/test_finetuning.py | 57 ++++++++++++++++++++++++++++----- 3 files changed, 101 insertions(+), 9 deletions(-) diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py index f6b39c928e..f6f5458c7d 100644 --- a/flash/core/finetuning.py +++ b/flash/core/finetuning.py @@ -13,7 +13,7 @@ # limitations under the License. import os from functools import partial -from typing import Iterable, Optional, Tuple, Union +from typing import Iterable, Optional, Tuple, Union, List, Dict, Any from pytorch_lightning import LightningModule from pytorch_lightning.callbacks import BaseFinetuning @@ -215,3 +215,53 @@ def __init__( train_bn: bool = True, ): super().__init__(FinetuningStrategies.UNFREEZE_MILESTONES, strategy_metadata, train_bn) + + +class FlashDeepSpeedFinetuning(FlashBaseFinetuning): + """FlashDeepSpeedFinetuning can be used to create a custom Flash Finetuning Callback which works with DeepSpeed. + DeepSpeed cannot store and load its parameters when working with pytorch-lightning. + So FlashDeepSpeedFinetuning override `_store` not to store its parameters.""" + + def _store( + self, + pl_module: "pl.LightningModule", + opt_idx: int, + num_param_groups: int, + current_param_groups: List[Dict[str, Any]], + ) -> None: + pass + + +class NoFreezeDeepSpeed(FlashDeepSpeedFinetuning): + def __init__(self, train_bn: bool = True): + super().__init__(FinetuningStrategies.NO_FREEZE, train_bn) + + +class FreezeDeepSpeed(FlashDeepSpeedFinetuning): + def __init__(self, train_bn: bool = True): + super().__init__(FinetuningStrategies.FREEZE, train_bn) + + +class FreezeUnfreezeDeepSpeed(FlashDeepSpeedFinetuning): + def __init__( + self, + strategy_metadata: int, + train_bn: bool = True, + ): + super().__init__(FinetuningStrategies.FREEZE_UNFREEZE, strategy_metadata, train_bn) + + +class UnfreezeMilestonesDeepSpeed(FlashDeepSpeedFinetuning): + def __init__( + self, + strategy_metadata: Tuple[Tuple[int, int], int], + train_bn: bool = True, + ): + super().__init__(FinetuningStrategies.UNFREEZE_MILESTONES, strategy_metadata, train_bn) + + +for strategy in FinetuningStrategies: + _FINETUNING_STRATEGIES_REGISTRY( + name=f"{strategy.value}_deepspeed", + fn=partial(FlashDeepSpeedFinetuning, strategy_key=strategy), + ) diff --git a/flash/core/utilities/imports.py b/flash/core/utilities/imports.py index ce48639b71..1c1ce6fe8b 100644 --- a/flash/core/utilities/imports.py +++ b/flash/core/utilities/imports.py @@ -118,6 +118,7 @@ def _compare_version(package: str, op: Callable, version: str, use_base_version: _BAAL_AVAILABLE = _module_available("baal") _TORCH_OPTIMIZER_AVAILABLE = _module_available("torch_optimizer") _SENTENCE_TRANSFORMERS_AVAILABLE = _module_available("sentence_transformers") +_DEEPSPEED_AVAILABLE = _module_available("deepspeed") if _PIL_AVAILABLE: diff --git a/tests/core/test_finetuning.py b/tests/core/test_finetuning.py index 67c63647b0..f56cfa9d47 100644 --- a/tests/core/test_finetuning.py +++ b/tests/core/test_finetuning.py @@ -27,7 +27,8 @@ import flash from flash.core.model import Task -from flash.core.utilities.imports import _CORE_TESTING +from flash.core.finetuning import _FINETUNING_STRATEGIES_REGISTRY +from flash.core.utilities.imports import _CORE_TESTING, _DEEPSPEED_AVAILABLE from tests.helpers.boring_model import BoringModel @@ -138,17 +139,38 @@ def on_train_epoch_start(self, trainer, pl_module): @pytest.mark.skipif(not _CORE_TESTING, reason="Not testing core.") @pytest.mark.parametrize( - "strategy", + "strategy, plugins", [ - "no_freeze", - "freeze", - ("freeze_unfreeze", 1), - ("unfreeze_milestones", ((5, 10), 5)), + ("no_freeze", None), + ("freeze", None), + (("freeze_unfreeze", 1), None), + (("unfreeze_milestones", ((5, 10), 5)), None), + pytest.param( + "no_freeze_deepspeed", + "deepspeed", + marks=pytest.mark.skipif(not _DEEPSPEED_AVAILABLE, reason="DeepSpeed not installed"), + ), + pytest.param( + "freeze_deepspeed", + "deepspeed", + marks=pytest.mark.skipif(not _DEEPSPEED_AVAILABLE, reason="DeepSpeed not installed"), + ), + pytest.param( + ("freeze_unfreeze_deepspeed", 1), + "deepspeed", + marks=pytest.mark.skipif(not _DEEPSPEED_AVAILABLE, reason="DeepSpeed not installed"), + ), + pytest.param( + ("unfreeze_milestones_deepspeed", ((5, 10), 5)), + "deepspeed", + marks=pytest.mark.skipif(not _DEEPSPEED_AVAILABLE, reason="DeepSpeed not installed"), + ), ], ) -def test_finetuning_with_none_return_type(strategy): +def test_finetuning_with_none_return_type(strategy, plugins): + gpus = 0 if plugins is None else 1 task = TestTaskWithoutFinetuning(loss_fn=F.nll_loss) - trainer = flash.Trainer(max_epochs=1, limit_train_batches=10) + trainer = flash.Trainer(max_epochs=1, limit_train_batches=10, gpus=gpus, plugins=plugins) ds = DummyDataset() trainer.finetune(task, train_dataloader=DataLoader(ds), strategy=strategy) @@ -223,3 +245,22 @@ def test_finetuning_errors_and_exceptions(strategy): ds = DummyDataset() with pytest.raises(MisconfigurationException): trainer.finetune(task, train_dataloader=DataLoader(ds), strategy=strategy) + + +@pytest.mark.parametrize( + "strategy_key, strategy_metadata", + [ + ("no_freeze", None), + ("freeze", None), + ("freeze_unfreeze", 2), + ("unfreeze_milestones", ((5, 10), 15)), + ], +) +def test_deepspeed_finetuning_strategy_key(strategy_key, strategy_metadata): + deepspeed_strategy_key = f"{strategy_key}_deepspeed" + + strategy = _FINETUNING_STRATEGIES_REGISTRY.get(key=strategy_key)(strategy_metadata=strategy_metadata).strategy + deepspeed_strategy = _FINETUNING_STRATEGIES_REGISTRY.get(key=deepspeed_strategy_key)( + strategy_metadata=strategy_metadata + ).strategy + assert strategy == deepspeed_strategy From 7e0f2e654ab7262e4b5026a92bdb9a5f06976c12 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 3 Jul 2022 05:53:43 +0000 Subject: [PATCH 2/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- flash/core/finetuning.py | 11 +++++++---- tests/core/test_finetuning.py | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py index f6f5458c7d..2102398bb5 100644 --- a/flash/core/finetuning.py +++ b/flash/core/finetuning.py @@ -13,7 +13,7 @@ # limitations under the License. import os from functools import partial -from typing import Iterable, Optional, Tuple, Union, List, Dict, Any +from typing import Any, Dict, Iterable, List, Optional, Tuple, Union from pytorch_lightning import LightningModule from pytorch_lightning.callbacks import BaseFinetuning @@ -218,9 +218,12 @@ def __init__( class FlashDeepSpeedFinetuning(FlashBaseFinetuning): - """FlashDeepSpeedFinetuning can be used to create a custom Flash Finetuning Callback which works with DeepSpeed. - DeepSpeed cannot store and load its parameters when working with pytorch-lightning. - So FlashDeepSpeedFinetuning override `_store` not to store its parameters.""" + """FlashDeepSpeedFinetuning can be used to create a custom Flash Finetuning Callback which works with + DeepSpeed. + + DeepSpeed cannot store and load its parameters when working with pytorch-lightning. So FlashDeepSpeedFinetuning + override `_store` not to store its parameters. + """ def _store( self, diff --git a/tests/core/test_finetuning.py b/tests/core/test_finetuning.py index f56cfa9d47..a37bc6e4fe 100644 --- a/tests/core/test_finetuning.py +++ b/tests/core/test_finetuning.py @@ -26,8 +26,8 @@ from torch.utils.data import DataLoader import flash -from flash.core.model import Task from flash.core.finetuning import _FINETUNING_STRATEGIES_REGISTRY +from flash.core.model import Task from flash.core.utilities.imports import _CORE_TESTING, _DEEPSPEED_AVAILABLE from tests.helpers.boring_model import BoringModel From 1413a1685767b6373c0a94fbf16b22a49b9468ef Mon Sep 17 00:00:00 2001 From: Masahiro Wada Date: Sun, 3 Jul 2022 06:21:27 +0000 Subject: [PATCH 3/9] Fix type --- flash/core/finetuning.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py index 2102398bb5..4a68073979 100644 --- a/flash/core/finetuning.py +++ b/flash/core/finetuning.py @@ -227,7 +227,7 @@ class FlashDeepSpeedFinetuning(FlashBaseFinetuning): def _store( self, - pl_module: "pl.LightningModule", + pl_module: LightningModule, opt_idx: int, num_param_groups: int, current_param_groups: List[Dict[str, Any]], From fdeb8fd192fc9a8ee71388d220b3f1e864b58243 Mon Sep 17 00:00:00 2001 From: Masahiro Wada Date: Sun, 3 Jul 2022 07:25:28 +0000 Subject: [PATCH 4/9] Add documents about finetunig with DeepSpeed --- docs/source/general/finetuning.rst | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/source/general/finetuning.rst b/docs/source/general/finetuning.rst index 42cb873e29..3c79b5f850 100644 --- a/docs/source/general/finetuning.rst +++ b/docs/source/general/finetuning.rst @@ -241,3 +241,13 @@ For even more customization, create your own finetuning callback. Learn more abo :hide: ... + +Working with DeepSpeed +====================== + +If you are using DeepSpeed, you can use the following strategies. The usage of the followings are the same as the above. But finetuning with DeepSpeed doesn't support load and store its parameters. + +* freeze_deepspeed +* no_freeze_deepspeed +* freeze_unfreeze_deepspeed +* unfreeze_milestones_deepspeed \ No newline at end of file From 7950b7ff2193ca72a2e21e218333706a7368b917 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 3 Jul 2022 07:26:18 +0000 Subject: [PATCH 5/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/source/general/finetuning.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/general/finetuning.rst b/docs/source/general/finetuning.rst index 3c79b5f850..908b94f25a 100644 --- a/docs/source/general/finetuning.rst +++ b/docs/source/general/finetuning.rst @@ -250,4 +250,4 @@ If you are using DeepSpeed, you can use the following strategies. The usage of t * freeze_deepspeed * no_freeze_deepspeed * freeze_unfreeze_deepspeed -* unfreeze_milestones_deepspeed \ No newline at end of file +* unfreeze_milestones_deepspeed From fd7177da40e0a746197438e6b0a067f6529fca4d Mon Sep 17 00:00:00 2001 From: Masahiro Wada Date: Fri, 22 Jul 2022 20:21:24 +0900 Subject: [PATCH 6/9] Update docs/source/general/finetuning.rst Co-authored-by: Kushashwa Ravi Shrimali --- docs/source/general/finetuning.rst | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/docs/source/general/finetuning.rst b/docs/source/general/finetuning.rst index 908b94f25a..f71ab35bdc 100644 --- a/docs/source/general/finetuning.rst +++ b/docs/source/general/finetuning.rst @@ -245,9 +245,9 @@ For even more customization, create your own finetuning callback. Learn more abo Working with DeepSpeed ====================== -If you are using DeepSpeed, you can use the following strategies. The usage of the followings are the same as the above. But finetuning with DeepSpeed doesn't support load and store its parameters. +If you are using DeepSpeed, you can use the following strategies. The usage of the following strategies is the same as listed above, but finetuning with DeepSpeed doesn't yet support the loading and storing of its parameters. -* freeze_deepspeed -* no_freeze_deepspeed -* freeze_unfreeze_deepspeed -* unfreeze_milestones_deepspeed +* ``freeze_deepspeed`` +* ``no_freeze_deepspeed`` +* ``freeze_unfreeze_deepspeed`` +* ``unfreeze_milestones_deepspeed`` From ea90fb441e1b39a289e2cfdc18880e5748fc9cb8 Mon Sep 17 00:00:00 2001 From: Masahiro Wada Date: Fri, 22 Jul 2022 20:22:12 +0900 Subject: [PATCH 7/9] Update flash/core/finetuning.py Co-authored-by: Kushashwa Ravi Shrimali --- flash/core/finetuning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py index 4a68073979..5fc29e7ed3 100644 --- a/flash/core/finetuning.py +++ b/flash/core/finetuning.py @@ -221,8 +221,8 @@ class FlashDeepSpeedFinetuning(FlashBaseFinetuning): """FlashDeepSpeedFinetuning can be used to create a custom Flash Finetuning Callback which works with DeepSpeed. - DeepSpeed cannot store and load its parameters when working with pytorch-lightning. So FlashDeepSpeedFinetuning - override `_store` not to store its parameters. + DeepSpeed cannot store and load its parameters when working with Lightning. So FlashDeepSpeedFinetuning + overrides `_store` to not store its parameters. """ def _store( From 9541b38edb2b1d9459ed6d72b20cb81448742b23 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 22 Jul 2022 11:22:54 +0000 Subject: [PATCH 8/9] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- flash/core/finetuning.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py index 5fc29e7ed3..356b99f04f 100644 --- a/flash/core/finetuning.py +++ b/flash/core/finetuning.py @@ -221,8 +221,8 @@ class FlashDeepSpeedFinetuning(FlashBaseFinetuning): """FlashDeepSpeedFinetuning can be used to create a custom Flash Finetuning Callback which works with DeepSpeed. - DeepSpeed cannot store and load its parameters when working with Lightning. So FlashDeepSpeedFinetuning - overrides `_store` to not store its parameters. + DeepSpeed cannot store and load its parameters when working with Lightning. So FlashDeepSpeedFinetuning overrides + `_store` to not store its parameters. """ def _store( From 2b98655c0e1a12a91475a5e3d24e8f8b4e9f3f46 Mon Sep 17 00:00:00 2001 From: Kushashwa Ravi Shrimali Date: Tue, 2 Aug 2022 10:06:22 +0530 Subject: [PATCH 9/9] Add changelog entry --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 804b60dd19..9de75f4ad7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ### Added +- Added fine tuning strategies for DeepSpeed (with parameter loading and storing omitted) ([#1377](https://github.com/Lightning-AI/lightning-flash/pull/1377)) + - Added `figsize` and `limit_nb_samples` for showing batch images ([#1381](https://github.com/Lightning-AI/lightning-flash/pull/1381)) - Added support for `from_lists` for Tabular Classification and Regression ([#1337](https://github.com/PyTorchLightning/lightning-flash/pull/1337))