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,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import deepspeed
+
+from unit.common import DistributedTest
+from unit.simple_model import *
+
+import pytest
+
+
+class TestSparseCheckpoint(DistributedTest):
+    world_size = 2
+
+    @pytest.mark.parametrize(["to_save_model_has_embedding", "to_save_model_sparse"], [
+        [False, False],
+        [True, False],
+        [True, True],
+    ])
+    @pytest.mark.parametrize(["destination_has_embedding", "destination_sparse"], [
+        [False, False],
+        [True, False],
+        [True, True],
+    ])
+    def test_non_strict_load_sparse(self, tmpdir, to_save_model_has_embedding, to_save_model_sparse,
+                                    destination_has_embedding, destination_sparse):
+
+        class ModelNoEmbedding(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.linear = torch.nn.Linear(3, 1)
+
+            def forward(self, x):
+                return self.linear(x)
+
+        class ModelEmbedding(torch.nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                self.emb = torch.nn.Embedding(10, 3)
+                self.linear = torch.nn.Linear(3, 1)
+
+            def forward(self, x, offsets):
+                return self.linear(self.emb(x, offsets))
+
+        if to_save_model_has_embedding:
+            model_to_save = ModelEmbedding()
+        else:
+            model_to_save = ModelNoEmbedding()
+        if destination_has_embedding:
+            model_destination = ModelEmbedding()
+        else:
+            model_destination = ModelNoEmbedding()
+
+        engine_to_save, _, _, _ = deepspeed.initialize(model=model_to_save,
+                                                       config={
+                                                           "train_batch_size": 2,
+                                                           "sparse_gradients": to_save_model_sparse
+                                                       })
+        engine_destination, _, _, _ = deepspeed.initialize(model=model_destination,
+                                                           config={
+                                                               "train_batch_size": 2,
+                                                               "sparse_gradients": destination_sparse
+                                                           })
+
+        save_folder = os.path.join(tmpdir, 'saved_checkpoint')
+        save_tag = '1'
+
+        engine_to_save.save_checkpoint(save_folder, tag=save_tag)
+
+        is_sparse_destination = isinstance(model_destination, ModelEmbedding) and destination_sparse
+        if isinstance(model_destination, ModelEmbedding) and model_destination.emb.sparse:
+            assert "emb.weight" in engine_destination.sparse_tensor_module_names
+        engine_destination.load_checkpoint(save_folder,
+                                           tag=save_tag,
+                                           load_module_strict=False,
+                                           load_optimizer_states=False,
+                                           load_lr_scheduler_states=False,
+                                           load_module_only=False)
+        if isinstance(model_destination, ModelEmbedding) and isinstance(model_to_save, ModelEmbedding):
+            assert engine_destination.sparse_tensor_module_names == engine_to_save.sparse_tensor_module_names
+        elif isinstance(model_destination, ModelEmbedding):
+            assert not is_sparse_destination or "emb.weight" in engine_destination.sparse_tensor_module_names
+        else:
+            assert len(engine_destination.sparse_tensor_module_names) == 0