[BUG] EmbeddingBag sparse gradients skipped without error or warning

[Dipet]

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 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

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.

[jeffra]

Thank you for reporting this issue @Dipet! I've just pushed #1418, can you give it a try? I've added a unit test based on your repro code above, thank you for including this.