Repository: microsoft/DeepSpeed

Issue:
[BUG] EmbeddingBag sparse gradients skipped without error or warning
**Describe the bug**
When you try to use `EmbeddingBag` with sparse gradients, they will be skipped silently and will not be used in `all_reduce` logic because [`split_half_float_double_csr`](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/runtime/engine.py#L79) doesn't check `torch.cuda.sparse.FloatTensor` type.
I think we need to change logic and not iterate over data types, but rather iterate over tensors and check that their data type in supported data types.

**To Reproduce**
```python
import copy
import torch
import typing
import torch.distributed
import deepspeed

params = {
    "train_batch_size": 2048,
    "train_micro_batch_size_per_gpu": 1024,
    "steps_per_print": 1,
    "sparse_gradients": True,
}


class Adam(torch.optim.Optimizer):
    def __init__(
        self,
        dense_params: list,
        sparse_params: list,
        lr: float = 1e-3,
        betas: typing.Tuple[float, float] = (0.9, 0.999),
        eps: float = 1e-8,
        weight_decay: float = 0,
        amsgrad: bool = False,
    ):
        super().__init__(
            dense_params + sparse_params, dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, amsgrad=amsgrad)
        )

        self.adam = torch.optim.Adam(dense_params, lr, betas, eps, weight_decay, amsgrad)
        self.adam_sparse = torch.optim.SparseAdam(sparse_params, lr, betas, eps)

    def state_dict(self) -> dict:
        result = super().state_dict()
        result["adam"] = self.adam.state_dict()
        result["adam_sparse"] = self.adam_sparse.state_dict()
        return result

    def load_state_dict(self, state_dict: dict) -> None:
        state_dict = copy.deepcopy(state_dict)
        self.adam.load_state_dict(state_dict.pop("adam"))
        self.adam_sparse.load_state_dict(state_dict.pop("adam_sparse"))
        super().load_state_dict(state_dict)

    @torch.no_grad()
    def step(self, closure=None) -> typing.Optional[float]:
        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 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))


model = Model()
model, optim, _, _ = deepspeed.initialize(
    model=model, optimizer=Adam(list(model.linear.parameters()), list(model.emb.parameters())), config_params=params
)
loss = torch.nn.BCEWithLogitsLoss()

device = f"cuda:{torch.distributed.get_rank()}"
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)
model(x, offsets)
y = torch.tensor([[1.0], [0.0]], device=device)

model.train()
res = model(x, offsets)
model.backward(loss(res, y))
model.step()

```

**Expected behavior**
Throw an error if it encounters gradients with unsupported data types. 



Discussion:

Resolving patch:
diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py
--- a/deepspeed/runtime/engine.py
+++ b/deepspeed/runtime/engine.py
@@ -78,14 +78,18 @@
 
 
 def split_half_float_double_csr(tensors):
-    dtypes = [
+    supported_types = [
         "torch.cuda.HalfTensor",
         "torch.cuda.FloatTensor",
         "torch.cuda.DoubleTensor",
         CSRTensor.type()
     ]
+
+    for t in tensors:
+        assert t.type() in supported_types, f"attempting to reduce an unsupported grad type: {t.type()}"
+
     buckets = []
-    for i, dtype in enumerate(dtypes):
+    for i, dtype in enumerate(supported_types):
         bucket = [t for t in tensors if t.type() == dtype]
         if bucket:
             buckets.append((dtype, bucket))
@@ -1761,7 +1765,7 @@ def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000)
                         grads.append(grad_data)
 
         split_buckets = split_half_float_double_csr(grads)
-        for i, bucket_tuple in enumerate(split_buckets):
+        for _, bucket_tuple in enumerate(split_buckets):
             bucket_type, bucket = bucket_tuple
 
             if self.pipeline_parallelism:
@@ -1770,22 +1774,17 @@ def buffered_allreduce_fallback(self, grads=None, elements_per_buffer=500000000)
                 dp_group = groups.get_data_parallel_group()
 
             if bucket_type == CSRTensor.type():
-                # TODO: do we have to do something here?
                 self.csr_allreduce_no_retain(bucket, dp_group=dp_group)
-                #groups.get_data_parallel_group() if self.pipeline_parallelism else self.mpu.get_data_parallel_group())
             else:
-                self.allreduce_no_retain(
-                    bucket,
-                    dp_group=dp_group,
-                    #groups.get_data_parallel_group(),
-                    numel_per_bucket=elements_per_buffer)
+                self.allreduce_no_retain(bucket,
+                                         dp_group=dp_group,
+                                         numel_per_bucket=elements_per_buffer)
 
         if self.has_moe_layers:
             expert_split_buckets = split_half_float_double_csr(expert_grads)
             for i, bucket_tuple in enumerate(expert_split_buckets):
                 bucket_type, bucket = bucket_tuple
                 if bucket_type == CSRTensor.type():
-                    # TODO: do we have to do something here?
                     self.csr_allreduce_no_retain(bucket,
                                                  groups.get_expert_data_parallel_group())
                 else:


Test patch:
diff --git a/tests/unit/common.py b/tests/unit/common.py
--- a/tests/unit/common.py
+++ b/tests/unit/common.py
@@ -42,6 +42,9 @@ def dist_init(local_rank, num_procs, *func_args, **func_kwargs):
             os.environ['RANK'] = str(local_rank)
             os.environ['WORLD_SIZE'] = str(num_procs)
 
+            # turn off NCCL logging if set
+            os.environ.pop('NCCL_DEBUG', None)
+
             deepspeed.init_distributed(dist_backend=backend)
 
             if torch.cuda.is_available():
diff --git a/tests/unit/test_sparse_grads.py b/tests/unit/test_sparse_grads.py
new file mode 100644
--- /dev/null
+++ b/tests/unit/test_sparse_grads.py
@@ -0,0 +1,61 @@
+import torch
+import torch.distributed as dist
+import deepspeed
+import pytest
+from common import distributed_test
+
+
+def test_sparse_adam(tmpdir):
+    config_dict = {"train_batch_size": 2, "steps_per_print": 1, "sparse_gradients": 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
+
+    model = Model()
+    optimizer = Adam(list(model.linear.parameters()), list(model.emb.parameters()))
+
+    @distributed_test(world_size=[2])
+    def _test(model, optimizer):
+        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)
+        with pytest.raises(AssertionError):
+            engine.backward(loss(res, y))
+        engine.step()
+
+    _test(model, optimizer)