Repository: microsoft/DeepSpeed

Issue:
[BUG] DeepSpeed ZeRO INIT with Stage 3 is failing with device mismatch error in v0.9.3. Works properly with v0.9.2
**Describe the bug**
DeepSpeed ZeRO INIT with Stage 3 is failing with device mismatch error

**To Reproduce**
Steps to reproduce the behavior:
Run below command:
```
accelerate launch --num_processes=2 --num_machines=1 --machine_rank=0 --mixed_precision=fp16 --use_deepspeed --gradient_accumulation_steps=1 --gradient_clipping=1 --zero3_init_flag=True --zero3_save_16bit_model=True --zero_stage=3 --offload_optimizer_device=none --offload_param_device=none accelerate/src/accelerate/test_utils/scripts/external_deps/test_peak_memory_usage.py --output_dir=/tmp/tmpjw8nujjd --peak_memory_upper_bound=2800 --n_train=160 --n_val=160
```
See the failing slow tests here: https://github.com/huggingface/accelerate/actions/runs/5172590339/jobs/9317076815

**Expected behavior**
Normal training. 

**ds_report output**
```
Setting ds_accelerator to cuda (auto detect)
--------------------------------------------------
DeepSpeed C++/CUDA extension op report
--------------------------------------------------
NOTE: Ops not installed will be just-in-time (JIT) compiled at
      runtime if needed. Op compatibility means that your system
      meet the required dependencies to JIT install the op.
--------------------------------------------------
JIT compiled ops requires ninja
ninja .................. [OKAY]
--------------------------------------------------
op name ................ installed .. compatible
--------------------------------------------------
 [WARNING]  async_io requires the dev libaio .so object and headers but these were not found.
 [WARNING]  async_io: please install the libaio-dev package with apt
 [WARNING]  If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
 [WARNING]  sparse_attn requires a torch version >= 1.5 and < 2.0 but detected 2.0
 [WARNING]  using untested triton version (2.0.0), only 1.0.0 is known to be compatible
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/torch']
torch version .................... 2.0.0
deepspeed install path ........... ['/home/sourab/miniconda3/envs/ml/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.9.3, unknown, unknown
torch cuda version ............... 11.7
torch hip version ................ None
nvcc version ..................... 11.7
deepspeed wheel compiled w. ...... torch 2.0, cuda 11.7
```

**Launcher context**
Accelerate Launcher



Discussion:

Resolving patch:
diff --git a/deepspeed/runtime/zero/partition_parameters.py b/deepspeed/runtime/zero/partition_parameters.py
--- a/deepspeed/runtime/zero/partition_parameters.py
+++ b/deepspeed/runtime/zero/partition_parameters.py
@@ -220,6 +220,8 @@ class ZeroParamStatus(Enum):
 _orig_torch_zeros = torch.zeros
 _orig_torch_ones = torch.ones
 _orig_torch_full = torch.full
+_orig_torch_arange = torch.arange
+_orig_torch_eye = torch.eye
 
 
 def zero_wrapper_for_fp_tensor_constructor(fn: Callable, target_fp_dtype: torch.dtype) -> Callable:
@@ -424,11 +426,7 @@ def _init_subclass(cls, **kwargs):
             torch.nn.modules.module.Module.__init_subclass__ = classmethod(_init_subclass)
             torch.nn.modules.module.Module.apply = apply_with_gather(torch.nn.modules.module.Module._old_apply)
 
-            torch.Tensor.__new__ = get_new_tensor_fn_for_dtype(self.dtype)
-            torch.empty = zero_wrapper_for_fp_tensor_constructor(_orig_torch_empty, self.dtype)
-            torch.zeros = zero_wrapper_for_fp_tensor_constructor(_orig_torch_zeros, self.dtype)
-            torch.ones = zero_wrapper_for_fp_tensor_constructor(_orig_torch_ones, self.dtype)
-            torch.full = zero_wrapper_for_fp_tensor_constructor(_orig_torch_full, self.dtype)
+            self._add_tensor_creation_wrappers()
 
             if self.mem_efficient_linear:
                 print_rank_0(
@@ -476,6 +474,24 @@ def _set_dtype(self, ds_config, dtype):
         else:
             self.dtype = dtype or torch.half
 
+    def _add_tensor_creation_wrappers(self):
+        torch.Tensor.__new__ = get_new_tensor_fn_for_dtype(self.dtype)
+        torch.empty = zero_wrapper_for_fp_tensor_constructor(_orig_torch_empty, self.dtype)
+        torch.zeros = zero_wrapper_for_fp_tensor_constructor(_orig_torch_zeros, self.dtype)
+        torch.ones = zero_wrapper_for_fp_tensor_constructor(_orig_torch_ones, self.dtype)
+        torch.full = zero_wrapper_for_fp_tensor_constructor(_orig_torch_full, self.dtype)
+        torch.arange = zero_wrapper_for_fp_tensor_constructor(_orig_torch_arange, self.dtype)
+        torch.eye = zero_wrapper_for_fp_tensor_constructor(_orig_torch_eye, self.dtype)
+
+    def _remove_tensor_creation_wrappers(self):
+        torch.Tensor.__new__ = torch.Tensor.__old_new__
+        torch.empty = _orig_torch_empty
+        torch.zeros = _orig_torch_zeros
+        torch.ones = _orig_torch_ones
+        torch.full = _orig_torch_full
+        torch.arange = _orig_torch_arange
+        torch.eye = _orig_torch_eye
+
     def remove_wrappers(self):
 
         def _disable_class(cls):
@@ -491,11 +507,7 @@ def _disable_class(cls):
             torch.nn.modules.module.Module.__init_subclass__ = torch.nn.modules.module.Module._old_init_subclass
             torch.nn.modules.module.Module.apply = torch.nn.modules.module.Module._old_apply
 
-            torch.Tensor.__new__ = torch.Tensor.__old_new__
-            torch.empty = _orig_torch_empty
-            torch.zeros = _orig_torch_zeros
-            torch.ones = _orig_torch_ones
-            torch.full = _orig_torch_full
+            self._remove_tensor_creation_wrappers()
 
             # un doing it here will undo it during training
             # if self.mem_efficient_linear:


Test patch: