Repository: microsoft/DeepSpeed

Issue:
deepspeed.zero.Init causes infinite recursion error
**Describe the bug**

When the `deepspeed.zero.Init` wraps not only the model but also the deepspeed.initialize call, a RecursionError is raised. 
This happens in deepspeed 0.6.5 but NOT in 0.6.4. It blocks the integration with [Lightning Lite](https://github.com/Lightning-AI/lightning/issues/13191) where we until now wrapped the entire run() method with the context. 


**To Reproduce**
```py

import argparse
import os

import deepspeed
import torch
import torch.nn as nn


class TheModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2, bias=False)


config = {
    "activation_checkpointing": {
        "contiguous_memory_optimization": False,
        "cpu_checkpointing": False,
        "partition_activations": False,
        "synchronize_checkpoint_boundary": False,
    },
    "aio": {
        "block_size": 1048576,
        "overlap_events": True,
        "queue_depth": 8,
        "single_submit": False,
        "thread_count": 1,
    },
    "train_micro_batch_size_per_gpu": 1,
    "zero_allow_untested_optimizer": True,
    "zero_optimization": {
        "allgather_bucket_size": 200000000,
        "allgather_partitions": True,
        "contiguous_gradients": True,
        "overlap_comm": True,
        "reduce_bucket_size": 200000000,
        "reduce_scatter": True,
        "stage": 3,
        "sub_group_size": 1000000000000,
    },
}


def worker(rank):
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "12234"
    os.environ["WORLD_SIZE"] = "2"
    os.environ["RANK"] = str(rank)
    os.environ["LOCAL_RANK"] = str(rank)
    deepspeed.init_distributed()

    model_parallel_context = deepspeed.zero.Init(
        remote_device="cpu", pin_memory=True, config_dict_or_path=config, dtype=torch.float32
    )

    # If the context goes over the model AND the deepspeed.initilize call, we get an infinite recursion error
    # This worked in 0.6.4, but not in 0.6.5
    with model_parallel_context:
        model = TheModel()

    # If the context only goes over the model, no error occurs (unindent the lines below)
        deepspeed_engine, deepspeed_optimizer, _, _ = deepspeed.initialize(
            args=argparse.Namespace(device_rank=rank),
            model=model,
            # model_parameters=model.parameters(),
            # optimizer=optimizer,
            dist_init_required=False,
            config=config,
        )


if __name__ == "__main__":
    torch.multiprocessing.spawn(worker, nprocs=2)

```

Output (may need to press ctrl+c on hang):

```
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 432, in __getattr__
    if name in dir(self):
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1847, in __dir__
    parameters = list(self._parameters.keys())
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.10/site-packages/deepspeed/runtime/engine.py", line 432, in __getattr__
    if name in dir(self):
  File "/home/adrian/anaconda3/envs/lightning/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1845, in __dir__
    module_attrs = dir(self.__class__)
RecursionError: maximum recursion depth exceeded while calling a Python object
```

**Expected behavior**
This worked in 0.6.4, so my assumption is that the change was unintentional. Git blame points to #1915. We weren't able to spot exactly which lines caused it, but suspect the getattr changes on the deepspeed engine.


**ds_report output**
```
--------------------------------------------------
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
--------------------------------------------------
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
 [WARNING]  please install triton==1.0.0 if you want to use sparse attention
sparse_attn ............ [NO] ....... [NO]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
 [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]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/adrian/anaconda3/envs/lightning/lib/python3.10/site-packages/torch']
torch version .................... 1.11.0
torch cuda version ............... 11.3
torch hip version ................ None
nvcc version ..................... 11.1
deepspeed install path ........... ['/home/adrian/anaconda3/envs/lightning/lib/python3.10/site-packages/deepspeed']
deepspeed info ................... 0.6.4, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.11, cuda 11.3
```

**System info (please complete the following information):**
 - OS: Ubuntu
 - GPU count and types: 2x RTX3090
 - Interconnects (if applicable) [e.g., two machines connected with 100 Gbps IB]
 - Python version: 3.9
 - Any other relevant info about your setup

**Launcher context**
I'm launching using torch.multiprocessing for simplicity in reproducing, but the bug is unrelated to how it is getting launched.

**Docker context**
No docker

**Additional context**
- 



Discussion:
@awaelchli, thanks for reporting this issue. `zero.Init()` is meant to wrap model construction only. Is there a reason why you are wrapping the entire execution, including `deepspeed.initialize()`?
We are providing a wrapper around the user code in [Lightning Lite](https://pytorch-lightning.readthedocs.io/en/stable/starter/lightning_lite.html). It looks something like this:
```py

class Lite(LightningLite):
    def run(self):
        # users code goes here (model, training loop, etc.)

       model = ...
       model, optimizer = self.setup(model, optimizer)

       # train model
         
if __name__ == "__main__":
    lite = Lite(accelerator="gpu", devices=2)
    lite.run()

```

We have a deepspeed integration. Without any major changes to their code, the user can turn deepspeed on by changing this:

```py
Lite(accelerator="gpu", devices=2, strategy=DeepspeedStrategy(stage=3, ...))
```

Internally, we wrap the run method with the `deepspeed.zero.Init` context manager so the user doesn't have to call it themselves. Then it would go over the entire user's code, including the self.setup() call where we initialize the deepspeed. 
Previously this worked as it wouldn't interfere with the `deepspeed.initialize`. 

It would be great if this could still be supported. If not, we'd need to introduce a context manager in Lite that the user has to call over their model instantiation. 

cc @carmocca
@awaelchli, thanks for sharing this context. Honestly, I don't know why it worked previously as wrapping `deepspeed.initialize()` was not an intended use for `zero.Init`. So, I am quite puzzled as to the recent changes that could have broken the integration. I will try to repro the failure. 

Resolving patch:
diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py
--- a/deepspeed/__init__.py
+++ b/deepspeed/__init__.py
@@ -115,6 +115,9 @@ def initialize(args=None,
         __git_branch__),
              ranks=[0])
 
+    # Disable zero.Init context if it's currently enabled
+    zero.partition_parameters.shutdown_init_context()
+
     assert model is not None, "deepspeed.initialize requires a model"
 
     if not isinstance(model, PipelineModule):
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
@@ -35,6 +35,7 @@
 
 param_count = 0
 partitioned_param_data_shape = [0]
+zero_init_enabled = False
 
 
 def _dist_allgather_fn(input_tensor: Tensor, output_tensor: Tensor, group=None):
@@ -258,8 +259,10 @@ def __init__(self,
         assert self.dtype in [torch.half, torch.bfloat16, torch.float], f"Invalid data type {self.dtype}, allowed values are [torch.half, torch.bfloat16, torch.float]"
 
     def __enter__(self):
+        global zero_init_enabled
         if not self.enabled:
             return
+        zero_init_enabled = True
 
         def apply_with_gather(orig_module_apply_fn: Callable) -> Callable:
             """many models make use of child modules like Linear or Embedding which
@@ -404,28 +407,7 @@ def __exit__(self, exc_type, exc_value, traceback):
         if not self.enabled:
             return
 
-        def _disable_class(cls):
-            cls.__init__ = cls._old_init
-
-        # Replace .__init__() for all existing subclasses of torch.nn.Module
-        for subclass in get_all_subclasses(torch.nn.modules.module.Module):
-            _disable_class(subclass)
-
-        # putting methods back the way we found them
-        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
-
-        # un doing it here will undo it during training
-        # if self.mem_efficient_linear:
-        #    torch.nn.functional.linear = self.linear_bk
-        #        if self.mem_efficient_linear:
-        #            torch.nn.functional.linear = self.linear_bk
+        shutdown_init_context()
 
         if dist.get_rank() == 0:
             logger.info("finished initializing model with %.2fB parameters",
@@ -454,6 +436,38 @@ def _set_dtype(self, ds_config, dtype):
             self.dtype = dtype or torch.half
 
 
+def shutdown_init_context():
+    global zero_init_enabled
+
+    if not zero_init_enabled:
+        return
+
+    def _disable_class(cls):
+        cls.__init__ = cls._old_init
+
+    # Replace .__init__() for all existing subclasses of torch.nn.Module
+    for subclass in get_all_subclasses(torch.nn.modules.module.Module):
+        _disable_class(subclass)
+
+    # putting methods back the way we found them
+    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
+
+    # un doing it here will undo it during training
+    # if self.mem_efficient_linear:
+    #    torch.nn.functional.linear = self.linear_bk
+    #        if self.mem_efficient_linear:
+    #            torch.nn.functional.linear = self.linear_bk
+
+    zero_init_enabled = False
+
+
 class AllGatherHandle:
     def __init__(self, handle, param: Parameter) -> None:
         if param.ds_status != ZeroParamStatus.INFLIGHT:
diff --git a/op_builder/transformer_inference.py b/op_builder/transformer_inference.py
--- a/op_builder/transformer_inference.py
+++ b/op_builder/transformer_inference.py
@@ -1,5 +1,4 @@
 import torch
-from packaging import version
 from .builder import CUDAOpBuilder, installed_cuda_version
 
 
@@ -18,7 +17,7 @@ def is_compatible(self, verbose=True):
         cuda_okay = True
         if not self.is_rocm_pytorch() and torch.cuda.is_available():
             sys_cuda_major, _ = installed_cuda_version()
-            torch_cuda_major = version.parse(torch.version.cuda).major
+            torch_cuda_major = int(torch.version.cuda.split('.')[0])
             cuda_capability = torch.cuda.get_device_properties(0).major
             if cuda_capability >= 8:
                 if torch_cuda_major < 11 or sys_cuda_major < 11:


Test patch:
diff --git a/tests/unit/test_zero_context.py b/tests/unit/test_zero_context.py
--- a/tests/unit/test_zero_context.py
+++ b/tests/unit/test_zero_context.py
@@ -360,3 +360,30 @@ def test_subclass_param_init():
         assert torch.equal(model.param, ones + 1)
         assert torch.equal(model.param_pa, ones + 2)
         assert torch.equal(model.param_grandpa, ones + 3)
+
+
+@distributed_test(world_size=2)
+def test_ds_init_w_zinit():
+    ds_config = {
+        "train_batch_size": 2,
+        "steps_per_print": 1,
+        "optimizer": {
+            "type": "Adam",
+            "params": {
+                "lr": 0.00015
+            }
+        }
+    }
+
+    class Model(torch.nn.Module):
+        def __init__(self):
+            super(Model, self).__init__()
+            self.linear = torch.nn.Linear(4, 4)
+
+        def magic(self):
+            return 42
+
+    with deepspeed.zero.Init():
+        model = Model()
+        engine, *_ = deepspeed.initialize(model=model, config=ds_config, model_parameters=model.parameters())
+    assert engine.magic() == 42