chore: import upstream snapshot with attribution

2026-07-13 13:18:33 +08:00
commit 4ececc111a
2017 changed files with 331736 additions and 0 deletions
@@ -0,0 +1,93 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+
+from unit.common import DistributedTest
+
+from transformers import GPT2Config, VisionEncoderDecoderConfig, VisionEncoderDecoderModel, ViTConfig
+from transformers.integrations.deepspeed import HfDeepSpeedConfig
+
+import deepspeed
+
+
+def _create_tiny_vision_encoder_decoder_model(model_path):
+    encoder_config = ViTConfig(image_size=8,
+                               patch_size=4,
+                               num_hidden_layers=1,
+                               hidden_size=8,
+                               num_attention_heads=2,
+                               intermediate_size=16)
+    decoder_config = GPT2Config(vocab_size=32,
+                                n_positions=16,
+                                n_embd=8,
+                                n_layer=1,
+                                n_head=2,
+                                bos_token_id=0,
+                                eos_token_id=1,
+                                add_cross_attention=True,
+                                is_decoder=True)
+    config = VisionEncoderDecoderConfig.from_encoder_decoder_configs(encoder_config, decoder_config)
+    model = VisionEncoderDecoderModel(config)
+    model.save_pretrained(model_path, safe_serialization=False)
+
+
+class TestNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_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):
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model = torch.nn.Linear(4, 4)
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model.weight, "ds_id")
+
+        deepspeed_engine, *_ = deepspeed.initialize(model=model, config_params=ds_config)
+
+
+class TestShutdownInNestingInit(DistributedTest):
+    world_size = 1
+
+    def test_shutdown_in_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):
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model1 = torch.nn.Linear(4, 4)
+
+            assert hasattr(model1.weight, "ds_id")
+            deepspeed_engine1, *_ = deepspeed.initialize(model=model1, config_params=ds_config)
+            with deepspeed.zero.Init(config_dict_or_path=ds_config):
+                model2 = torch.nn.Linear(4, 4)
+
+        # ensure that zero3 processed the parameter
+        assert hasattr(model2.weight, "ds_id")
+        deepspeed_engine2, *_ = deepspeed.initialize(model=model2, config_params=ds_config)
+
+
+class TestNestedParallelInit(DistributedTest):
+    world_size = 1
+
+    # Testing a model with composed and nested zero.Inits, with 3 zero.Init contexts, 1 parent and 2 children.
+    # The skeleton of the model is like so
+    #
+    # class VisionEncoderDecoderModel(...)::
+    #     def __init__(self):
+    #             encoder = AutoModel.from_config(config.encoder)
+    #             decoder = AutoModelForCausalLM.from_config(config.decoder)
+    #
+    # And the user calls like below:
+    # VisionEncoderDecoderModel.from_pretrained(...)
+    # which calls this constructor inside zero.Init
+
+    def test_nested_parallel_init(self, tmp_path):
+        ds_config = dict(train_batch_size=1, zero_optimization=dict(stage=3))
+        _create_tiny_vision_encoder_decoder_model(tmp_path)
+        dschf = HfDeepSpeedConfig(ds_config)  # keep this object alive
+        model = VisionEncoderDecoderModel.from_pretrained(str(tmp_path), local_files_only=True)
+        assert all([hasattr(p, 'ds_id') for p in model.parameters()])