chore: import upstream snapshot with attribution

tests/unit/runtime/sparse_tensor/test_averaging_sparse_gradients.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import pytest
+import deepspeed
+from unit.common import DistributedTest
+from unit.util import skip_on_arch
+from deepspeed.accelerator import get_accelerator
+
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True)
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.emb = torch.nn.EmbeddingBag(10, 3, mode="sum", sparse=True)
+        self.linear = torch.nn.Linear(3, 1)
+
+    def forward(self, x, offsets):
+        return self.linear(self.emb(x, offsets))
+
+
+class Adam(torch.optim.Optimizer):
+
+    def __init__(self, dense_params, sparse_params):
+        super().__init__(dense_params + sparse_params, defaults={})
+        self.adam = torch.optim.Adam(dense_params)
+        self.adam_sparse = torch.optim.SparseAdam(sparse_params)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss_1 = self.adam.step(closure)
+        loss_2 = self.adam_sparse.step(closure)
+
+        if loss_1 is not None and loss_2 is not None:
+            return loss_1 + loss_2
+        return loss_1 or loss_2
+
+
+def get_model_optimizer():
+    torch.manual_seed(0)
+    model = Model()
+    optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters()))
+    return model, optimizer
+
+
+def get_data(device):
+    x = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long, device=device)
+    offsets = torch.tensor([0, 4], dtype=torch.long, device=device)
+    y = torch.tensor([[1.0], [0.0]], device=device)
+    return x, offsets, y
+
+
+class TestSparseAdam(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        skip_on_arch(min_arch=7)
+
+        config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True}
+        model, optimizer = get_model_optimizer()
+        loss = torch.nn.BCEWithLogitsLoss()
+        engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict)
+
+        x, offsets, y = get_data(engine.device)
+
+        engine.gradient_average = True
+        res = engine(x, offsets)
+        engine.backward(loss(res, y))
+
+        averaged_grads = {}
+        for k, v in engine.named_parameters():
+            grad = v.grad.to_dense() if v.grad.is_sparse else v.grad
+            averaged_grads[k] = grad
+            v.grad = None
+
+        engine.gradient_average = False
+        res = engine(x, offsets)
+        engine.backward(loss(res, y))
+
+        for k, v in engine.named_parameters():
+            grad = v.grad.to_dense() if v.grad.is_sparse else v.grad
+            assert torch.allclose(grad, averaged_grads[k] * engine.world_size)

tests/unit/runtime/sparse_tensor/test_csr.py

@@ -0,0 +1,55 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import random
+from deepspeed.runtime.sparse_tensor import SparseTensor
+
+
+def test_csr_addition_self():
+    row_count = 10
+    random.seed(1234)
+
+    x = torch.ones(1, 5)
+    for i in range(row_count - 1):
+        if random.random() > 0.75:
+            x = torch.cat([x, torch.ones(1, 5)])
+        else:
+            x = torch.cat([x, torch.zeros(1, 5)])
+    dense_x = x.clone()
+    cx = SparseTensor(x)
+
+    assert torch.all(dense_x == cx.to_dense())
+
+    cx.add(cx)
+    assert torch.all(dense_x + dense_x == cx.to_dense())
+
+
+def test_csr_addition_different():
+    row_count = 10
+    random.seed(1234)
+
+    x = torch.ones(1, 5)
+    for i in range(row_count - 1):
+        if random.random() > 0.75:
+            x = torch.cat([x, torch.ones(1, 5)])
+        else:
+            x = torch.cat([x, torch.zeros(1, 5)])
+    dense_x = x.clone()
+    cx = SparseTensor(x)
+
+    y = torch.ones(1, 5)
+    for i in range(row_count - 1):
+        if random.random() > 0.75:
+            y = torch.cat([y, torch.ones(1, 5)])
+        else:
+            y = torch.cat([y, torch.zeros(1, 5)])
+    dense_y = y.clone()
+    cy = SparseTensor(y)
+
+    dense_sum = dense_x + dense_y
+    cx.add(cy)
+
+    assert torch.all(dense_sum == cx.to_dense())

tests/unit/runtime/sparse_tensor/test_sparse_grads.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import torch
+import pytest
+import deepspeed
+from unit.common import DistributedTest
+from deepspeed.accelerator import get_accelerator
+import deepspeed.utils.groups as groups
+
+if get_accelerator().device_name() == 'hpu':
+    pytest.skip("sparse_gradients not supported by HPU.", allow_module_level=True)
+
+
+class Model(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.emb = torch.nn.EmbeddingBag(10, 3, mode="sum", sparse=True)
+        self.linear = torch.nn.Linear(3, 1)
+
+    def forward(self, x, offsets):
+        return self.linear(self.emb(x, offsets))
+
+
+class Adam(torch.optim.Optimizer):
+
+    def __init__(self, dense_params, sparse_params):
+        super().__init__(dense_params + sparse_params, defaults={})
+        self.adam = torch.optim.Adam(dense_params)
+        self.adam_sparse = torch.optim.SparseAdam(sparse_params)
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss_1 = self.adam.step(closure)
+        loss_2 = self.adam_sparse.step(closure)
+
+        if loss_1 is not None and loss_2 is not None:
+            return loss_1 + loss_2
+        return loss_1 or loss_2
+
+
+class TestSparseAdam(DistributedTest):
+    world_size = 2
+
+    def test(self):
+        config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": True}
+
+        model = Model()
+        optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters()))
+        engine, _, _, _ = deepspeed.initialize(model=model, optimizer=optimizer, config=config_dict)
+        loss = torch.nn.BCEWithLogitsLoss()
+        x = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9], dtype=torch.long, device=engine.device)
+        offsets = torch.tensor([0, 4], dtype=torch.long, device=engine.device)
+        y = torch.tensor([[1.0], [0.0]], device=engine.device)
+        res = engine(x, offsets)
+        engine.backward(loss(res, y))
+        engine.step()
+
+        results = [engine.all_gather_scalar(i, groups._get_data_parallel_group()) for i in model.emb.parameters()]
+        for res in results:
+            assert torch.allclose(res[0], res[1])