Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
test_sparse.py 3.95 KB
Newer Older
aiss's avatar
push dsv0.8.2 version  
aiss committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96 
'''Copyright The Microsoft 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