wehub-resource-sync
89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
# 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
|