Repository: microsoft/DeepSpeed

Issue:
[Bug] RuntimeError: expected input to be on cuda
**Describe the bug**
When I try to run the Megatron deepspeed example(https://github.com/microsoft/Megatron-DeepSpeed/tree/main/examples) with latest master branch of deepspeed, I met following error:
![image](https://user-images.githubusercontent.com/13343117/183005422-b4b7890a-0bb2-47c8-b03b-572f7807211b.png)


**To Reproduce**
Steps to reproduce the behavior:
1. Go to '[Megatron Example](https://github.com/microsoft/Megatron-DeepSpeed/tree/main/examples)'
2. Install latest DeepSpeed(0.7.0+9dcfb93a) & Megatron(0.5.1)
3. Prepare the data: 
 - bash dataset/download_books.sh
 - bash dataset/download_vocab.sh
5. Start the training:
 bash examples/azure/run-benchmark-model.sh

The env I used is:
1 node/ 8GPUs, A100. 

**Expected behavior**
A clear and concise description of what you expected to happen.

**ds_report output**
Please run `ds_report` to give us details about your setup.

**Screenshots**
Error:
![image](https://user-images.githubusercontent.com/13343117/183005761-c8bf7957-3978-4519-a638-cb85ae004260.png)

DS_CONFIG:
![image](https://user-images.githubusercontent.com/13343117/183005828-05dccd62-96ee-440a-bb15-976b65bd10fa.png)


**System info (please complete the following information):**
 - ubuntu2004-cu115-py38-torch1110
 - 1 machines with x8 A100s each



Discussion:
I have the same issue as yours and I suggested a temporary fix on their Megatron-DeepSpeed github.
https://github.com/microsoft/Megatron-DeepSpeed/pull/71
> I have the same issue as yours and I suggested a temporary fix on their Megatron-DeepSpeed github. [microsoft/Megatron-DeepSpeed#71](https://github.com/microsoft/Megatron-DeepSpeed/pull/71)

Thanks @hibagus. Good catch here. Also I got the fix from @awan-10 where we can change the `CPU_OPTIM` to `--cpu-optimizer` from `""`. It would be like for the [file ](https://github.com/microsoft/Megatron-DeepSpeed/blob/main/examples/azure/run-benchmark-model.sh#:~:text=%23%20Set%20to,CPU_OPTIM%3D%22%20%22) 
```sh
OFFLOAD_DEVICE="none"
CPU_OPTIM=" --cpu-optimizer"
```

Resolving patch:
diff --git a/deepspeed/runtime/zero/config.py b/deepspeed/runtime/zero/config.py
--- a/deepspeed/runtime/zero/config.py
+++ b/deepspeed/runtime/zero/config.py
@@ -9,7 +9,7 @@
 from enum import Enum
 from deepspeed.runtime.config_utils import get_scalar_param, DeepSpeedConfigModel
 from deepspeed.utils import logger
-from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig
+from .offload_config import DeepSpeedZeroOffloadParamConfig, DeepSpeedZeroOffloadOptimizerConfig, OffloadDeviceEnum
 
 # ZeRO optimization. By default, this optimization is not enabled.
 # Users have to configure the desired optimization (0 means disabled) in params.json as below example:
@@ -94,7 +94,9 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
         None,
         deprecated=True,
         new_param="offload_param",
-        new_param_fn=(lambda val: DeepSpeedZeroOffloadParamConfig() if val else None),
+        new_param_fn=(
+            lambda val: DeepSpeedZeroOffloadParamConfig(device=OffloadDeviceEnum.cpu)
+            if val else None),
     )
     cpu_offload_use_pin_memory: bool = Field(
         None,
@@ -106,8 +108,9 @@ class DeepSpeedZeroConfig(DeepSpeedConfigModel):
         None,
         deprecated=True,
         new_param="offload_optimizer",
-        new_param_fn=(lambda val: DeepSpeedZeroOffloadOptimizerConfig()
-                      if val else None),
+        new_param_fn=(
+            lambda val: DeepSpeedZeroOffloadOptimizerConfig(device=OffloadDeviceEnum.cpu)
+            if val else None),
     )
 
     # Stage3 Specific Parameters
diff --git a/deepspeed/runtime/zero/parameter_offload.py b/deepspeed/runtime/zero/parameter_offload.py
--- a/deepspeed/runtime/zero/parameter_offload.py
+++ b/deepspeed/runtime/zero/parameter_offload.py
@@ -187,7 +187,8 @@ def __init__(self,
         self.dtype = list(module.parameters())[0].dtype
         self.offload_device = None
         self.offload_param_pin_memory = False
-        if offload_param_config is not None:
+
+        if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none:
             self.offload_device = offload_param_config.device
             self.offload_param_pin_memory = offload_param_config.pin_memory
 
diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py
--- a/deepspeed/runtime/zero/stage3.py
+++ b/deepspeed/runtime/zero/stage3.py
@@ -446,14 +446,14 @@ def _get_param_coordinator(self, training):
 
     def _configure_offloading(self, offload_optimizer_config, offload_param_config):
         ###################### offload optimizer setup ##################################
-        if offload_optimizer_config is not None:
+        if offload_optimizer_config is not None and offload_optimizer_config.device != OffloadDeviceEnum.none:
             self.offload_optimizer = True
             self.offload_optimizer_pin_memory = offload_optimizer_config.pin_memory
             self.swap_optimizer = offload_optimizer_config.device == OffloadDeviceEnum.nvme
             self.offload_optimizer_fast_init = offload_optimizer_config.fast_init
 
         ###################### offload param setup ##################################
-        if offload_param_config is not None:
+        if offload_param_config is not None and offload_param_config.device != OffloadDeviceEnum.none:
             self.offload_param = True
             self.offload_param_pin_memory = offload_param_config.pin_memory
             self.params_in_nvme_and_cpu = offload_param_config.device == OffloadDeviceEnum.nvme


Test patch: