Repository: pytorch
Issue: Gloo errors when process's batch only indexes padding_idx of sparse embedding
Labels: oncall: distributed

### 🐛 Describe the bug

When a model uses embeddings that are sparse and padded (using padded_idx). Multi-process CPU training via DDP will produce Gloo errors of the form `op.preamble.length <= op.nbytes. 20 vs 16` any time a process's entire batch indexes the padded index.

I suspect this causes problems during gradient reduction because the process in question has an empty gradient tensor while the other processes have non-empty gradient tensors.

The error doesn't occur when embeddings are dense, unpadded, or when the entire batch indexes a single unpadded index.

## Repro



```python
import pytorch_lightning as pl
import torch

class TestModule(pl.LightningModule):
    def __init__(self, emb_size, padded, sparse, offset):
        super().__init__()
        self.offset = offset
        self.embedding = torch.nn.Embedding(
            emb_size, 1, padding_idx=0 if padded else None, sparse=sparse
        )
        self.criterion = torch.nn.MSELoss()

    def forward(self, x):
        # simulate a batch that only indexes one element in the embedding
        rnk = torch.distributed.get_rank()
        a = self.embedding(torch.ones(x[:, 0].size()).int() * (rnk + self.offset))
        return a

    def training_step(self, batch, batch_idx):
        x, y = batch
        x_hat = self.forward(x)
        loss = self.criterion(x_hat.flatten(), y.flatten())
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adagrad(
            self.parameters(),
            lr=0.5,
        )
        return optimizer


def train(padded: bool, sparse: bool, offset: int):
    emb_size = 12
    data_size = 1000

    ds = torch.utils.data.TensorDataset(
        torch.randint(0, emb_size, [data_size, 1]),
        torch.rand([data_size, 1]),
    )

    dl = torch.utils.data.DataLoader(ds, batch_size=10)

    model = TestModule(emb_size, padded, sparse, offset)
    trainer = pl.Trainer(
        num_nodes=3,
        devices=4,
        max_epochs=3,
        accelerator="cpu",
    )
    trainer.fit(model, dl)


train(True, True, 0)   # fails

train(False, True, 0)   # not padded, succeeds

train(True, False, 0)   # not sparse, succeeds

train(True, True, 1)   # doesn't index the padded_idx, succeeds
```

## Error message
```
E1020 12:56:59.375689  1777 ExceptionTracer.cpp:222] exception stack complete
terminate called after throwing an instance of 'gloo::EnforceNotMet'
  what():  [enforce fail at gloo/transport/tcp/pair.cc:510] op.preamble.length <= op.nbytes. 16 vs 8
*** Aborted at 1666295819 (Unix time, try 'date -d @1666295819') ***
*** Signal 6 (SIGABRT) (0x8068200000507) received by PID 1287 (pthread TID 0x7f6b53e00640) (linux TID 1777) (maybe from PID 1287, UID 525954) (code: -6), stack trace: ***
    @ 000000002f76a238 folly::symbolizer::(anonymous namespace)::signalHandler(int, siginfo_t*, void*)
    @ 0000000000000000 (unknown)
    @ 000000000009c9d3 __GI___pthread_kill
    @ 00000000000444ec __GI_raise
    @ 000000000002c432 __GI_abort
    @ 00000000000a3fd4 __gnu_cxx::__verbose_terminate_handler()
    @ 0000000016af6900 aiplatform::error_reporting::(anonymous namespace)::logExceptionDetailsAndExit(folly::exception_wrapper const&, bool, std::optional<int>, aiplatform::error_reporting::TerminationLogger const*)
    @ 0000000016af25f7 aiplatform::error_reporting::(anonymous namespace)::terminateHandler()
    @ 00000000000a1b39 __cxxabiv1::__terminate(void (*)())
    @ 00000000000a1ba4 std::terminate()
    @ 00000000000a1e6f __cxa_throw
    @ 0000000000003c6e __cxa_throw
    @ 0000000031b7a283 gloo::transport::tcp::Pair::prepareRead(gloo::transport::tcp::Op&, gloo::NonOwningPtr<gloo::transport::tcp::UnboundBuffer>&, iovec&)
    @ 0000000031b7a398 gloo::transport::tcp::Pair::read()
    @ 0000000031b7b728 gloo::transport::tcp::Pair::handleReadWrite(int)
    @ 0000000031b7f7c3 gloo::transport::tcp::Pair::handleEvents(int)
    @ 0000000031b7622d gloo::transport::tcp::Loop::run()
    @ 00000000000df4e4 execute_native_thread_routine
    @ 000000000009ac0e start_thread
    @ 000000000012d1db __clone3
Fatal Python error: Aborted
```

### Versions

collect_env.py failed, but I'm running on trunk

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu