chore: import upstream snapshot with attribution

tests/unit/runtime/zero/test_zero_dynamic_class.py

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from unit.common import DistributedTest
+
+import deepspeed
+
+
+class TestNewClassDeclaredNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_new_class_declared_nesting_init(self):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+
+        with deepspeed.zero.Init(config_dict_or_path=ds_config):
+
+            class MyModel(torch.nn.Module):
+
+                def __init__(self):
+                    super().__init__()
+                    self.fc = torch.nn.Linear(4, 4)
+
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model = MyModel()
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model.fc.weight, "ds_id")
+        deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)
+
+
+class TestNewClassDeclaredInsideNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_new_class_declared_inside_nesting_init(self):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+
+        with deepspeed.zero.Init(config_dict_or_path=ds_config):
+
+            class MyModel(torch.nn.Module):
+
+                def __init__(self):
+                    super().__init__()
+                    self.fc = torch.nn.Linear(1, 1)
+
+            model = MyModel()
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model.fc.weight, "ds_id")
+        deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)