Repository: Lightning-Universe/lightning-flash

Issue:
Flash DeepSpeedPlugin error
## 🐛 Bug
Hi, I'm completely new Flash, I'm having this error when trying to use DeepSpeedPlugin plugin on the Flash trainer.
The code is 

### To Reproduce

```python

import torch

import flash
from flash.audio import SpeechRecognition, SpeechRecognitionData
from flash.core.data.utils import download_data

# 1. Create the DataModule
download_data("https://pl-flash-data.s3.amazonaws.com/timit_data.zip", "./data")

datamodule = SpeechRecognitionData.from_json(
    "file",
    "text",
    train_file="data/timit/train.json",
    test_file="data/timit/test.json",
    batch_size=4,
)

# 2. Build the task
model = SpeechRecognition(backbone="facebook/wav2vec2-base-960h")

# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=1, gpus=torch.cuda.device_count(),plugins='deepspeed')
trainer.finetune(model, datamodule=datamodule, strategy="no_freeze")

# 4. Predict on audio files!
datamodule = SpeechRecognitionData.from_files(predict_files=["data/timit/example.wav"], batch_size=4)
predictions = trainer.predict(model, datamodule=datamodule)
print(predictions)

# 5. Save the model!
trainer.save_checkpoint("speech_recognition_model.pt")
```

I'm getting this
<img width="650" alt="Screenshot 2022-03-25 at 16 45 29" src="https://user-images.githubusercontent.com/25756717/160154674-653c8785-9633-466c-9325-8318ce3810ca.png">

I thing there is an issue with named_params somewhere 



Discussion:
Hey @Moumeneb1 Thanks for reporting this! We'll look into it and get back to you 😃
Is there anyone to look into it? If no one, may I try to look into it?
> Is there anyone to look into it? If no one, may I try to look into it?

Hi, @ar90n - Thank you for showing interest in this issue! It would be great if you want to take a look into it, as far as I'm aware - no one is actively working on this issue, so please feel free to pick it up. In case you face any issues, please feel free to reach out to me or @ethanwharris on our [PyTorch Lightning Slack](https://www.pytorchlightning.ai/community) workspace.

I've also assigned this issue to you! :)
@krshrimali 
Thank for your kindness. I'm going to get to this issue.
Hi, after I looked into this issue, I found the cause of it. It doesn’t depend on lightning-flash. It occurs with only pytorch-lightning. And there is some issue like this in pytorch-ligntning Github project page.  Therefore I added some comments on them and create a new issue about integration with pytorch-lightning and DeepSpeed. They are the followings.

* https://github.com/PyTorchLightning/pytorch-lightning/issues/13085
* https://github.com/PyTorchLightning/pytorch-lightning/issues/9833
* https://github.com/PyTorchLightning/pytorch-lightning/issues/12130

I think that if the above issue is solved, this issue will be solved automatically. I continue to try to solve them.
Hi, after my investigation, I found that finetuning using pytorch-lightning can not work with DeepSpeed correctly. Concretely, loading its parameters doesn't work when resuming training. 
As I commented in [this issue](https://github.com/PyTorchLightning/pytorch-lightning/issues/9833),  this is caused by pytorch-lightinig’s lifecycle which is the ordering of calling callbacks. So I think it is difficult to change it by the risk of the compatibility of older versions. But if we can omit the support of resuming training, there are simple workarounds as followings. This doesn’t support resuming but works.

```python
import torch
from functools import partial

import flash
from flash.audio import SpeechRecognition, SpeechRecognitionData
from flash.core.data.utils import download_data

from flash.core.finetuning import FlashBaseFinetuning, FinetuningStrategies, _FINETUNING_STRATEGIES_REGISTRY


class FlashDeepSpeedFinetuning(FlashBaseFinetuning):
    def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
        from pytorch_lightning.loops.utilities import _get_active_optimizers

        for opt_idx, optimizer in _get_active_optimizers(trainer.optimizers, trainer.optimizer_frequencies, 0):
            self.finetune_function(pl_module, trainer.current_epoch, optimizer, opt_idx)


class NoFreeze(FlashDeepSpeedFinetuning):
    def __init__(self, train_bn: bool = True):
        super().__init__(FinetuningStrategies.NO_FREEZE, train_bn)


_FINETUNING_STRATEGIES_REGISTRY(
    name="no_freeze_deepspeed",
    fn=partial(FlashDeepSpeedFinetuning, strategy_key=FinetuningStrategies.NO_FREEZE),
)


# 1. Create the DataModule
download_data("https://pl-flash-data.s3.amazonaws.com/timit_data.zip", "./data")

datamodule = SpeechRecognitionData.from_json(
    "file",
    "text",
    train_file="data/timit/train.json",
    test_file="data/timit/test.json",
    batch_size=4,
)

# 2. Build the task
model = SpeechRecognition(backbone="facebook/wav2vec2-base-960h")

# 3. Create the trainer and finetune the model
trainer = flash.Trainer(max_epochs=2, gpus=torch.cuda.device_count(), strategy="deepspeed")
trainer.finetune(model, datamodule=datamodule, strategy="no_freeze_deepspeed")

# 4. Predict on audio files!
datamodule = SpeechRecognitionData.from_files(predict_files=["data/timit/example.wav"], batch_size=4)
predictions = trainer.predict(model, datamodule=datamodule)
print(predictions)

# 5. Save the model!
torch.save(model.state_dict(), "./speech_recognition_model.pth")
```

What do you think about this workaround?
Hi, @ar90n - I just wanted to thank you for your continuous efforts on this issue! Apologies for not being able to respond on time.

Just to acknowledge that we are aware of this, I'll get back to you in a day or two. :) I appreciate the patience 🚀 and your efforts! ❤️ 
This is awesome @ar90n !!! Thanks for digging in to this. We would definitely welcome a PR to add your `no_freeze_deepspeed` strategy to Flash if that's something you're interested in? Also, do you know if it would be possible to make a `freeze` strategy to work with deepspeed too? That would unlock finetuning with deepspeed which would be quite cool :smiley:
@ethanwharris 
Thanks for your suggestion! I'm so interested in creating a PR about the `no_freeze_deepspeed` strategy.  And the `freeze` strategy doesn't work because of the same reason as `no_freeze`. So this is solved in the same way as the `no_freeze` strategy. I will create a PR about adding `no_freeze_deepspeed` and `freeze_deepspeed` strategies.


Resolving patch:
diff --git a/flash/core/finetuning.py b/flash/core/finetuning.py
--- 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 Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from pytorch_lightning import LightningModule
 from pytorch_lightning.callbacks import BaseFinetuning
@@ -215,3 +215,56 @@ 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 Lightning. So FlashDeepSpeedFinetuning overrides
+    `_store` to not store its parameters.
+    """
+
+    def _store(
+        self,
+        pl_module: 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
--- 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:


Test patch:
diff --git a/tests/core/test_finetuning.py b/tests/core/test_finetuning.py
--- a/tests/core/test_finetuning.py
+++ b/tests/core/test_finetuning.py
@@ -26,8 +26,9 @@
 from torch.utils.data import DataLoader
 
 import flash
+from flash.core.finetuning import _FINETUNING_STRATEGIES_REGISTRY
 from flash.core.model import Task
-from flash.core.utilities.imports import _CORE_TESTING
+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