Repository: pytorch
Issue: nn.Embedding weights are not synced across processes with DistributedDataParallel when other parameters are present
Labels: oncall: distributed,module: ddp

### 🐛 Describe the bug


If our model only contains a single `nn.Embedding` and no other parameters, then running DistributedDataParallel on two RTX 3090s leads to the expected behavior: the embeddings are the same across both processes. However, adding a `nn.Linear` (which isn't used at all) leads to the embeddings only being locally updated on each process. 

```python
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.multiprocessing as mp
import torch.distributed as dist


class TestModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.embedding = nn.Embedding(num_embeddings=2, embedding_dim=3)
        self.embedding.weight.data.fill_(0)
        # self.fc = nn.Linear(1, 1)  # Note: uncommenting this results in the bug

    def forward(self, idx):
        return self.embedding(idx).mean(dim=-1)


def main_worker(rank):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'
    dist.init_process_group("gloo", rank=rank, world_size=2)

    model = TestModel().to(rank)
    model = torch.nn.parallel.DistributedDataParallel(model)
    optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

    # dataset
    ids = torch.arange(2)
    targets = torch.arange(2).float() + 1
    dataset = torch.utils.data.TensorDataset(ids, targets)
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset, shuffle=True)
    train_loader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False,
                                               sampler=train_sampler, drop_last=True)

    print(f"Rank={rank} before update:", rank, model.module.embedding.weight)
    # 1 epoch of training
    for (id, target) in train_loader:
        preds = model(id.to(rank))
        optimizer.zero_grad()
        loss = F.mse_loss(preds, target.to(rank))
        loss.backward()
        optimizer.step()

    print(f"Rank={rank} after update:", rank, model.module.embedding.weight)
    dist.destroy_process_group()


if __name__ == '__main__':
    mp.spawn(main_worker, nprocs=2)

```

Output with just `nn.Embedding`, which matches expected behavior: 
```
Rank=1 before update: 1 Parameter containing:
tensor([[0., 0., 0.],
        [0., 0., 0.]], device='cuda:1', requires_grad=True)
Rank=0 before update: 0 Parameter containing:
tensor([[0., 0., 0.],
        [0., 0., 0.]], device='cuda:0', requires_grad=True)
Rank=1 after update: 1 Parameter containing:
tensor([[0.0333, 0.0333, 0.0333],
        [0.0667, 0.0667, 0.0667]], device='cuda:1', requires_grad=True)
Rank=0 after update: 0 Parameter containing:
tensor([[0.0333, 0.0333, 0.0333],
        [0.0667, 0.0667, 0.0667]], device='cuda:0', requires_grad=True)
```

Output when other parameters are present (embedding weights are not synced): 
```
Rank=1 before update: 1 Parameter containing:
tensor([[0., 0., 0.],
        [0., 0., 0.]], device='cuda:1', requires_grad=True)
Rank=1 after update: 1 Parameter containing:
tensor([[0.0000, 0.0000, 0.0000],
        [0.1333, 0.1333, 0.1333]], device='cuda:1', requires_grad=True)
Rank=0 before update: 0 Parameter containing:
tensor([[0., 0., 0.],
        [0., 0., 0.]], device='cuda:0', requires_grad=True)
Rank=0 after update: 0 Parameter containing:
tensor([[0.0667, 0.0667, 0.0667],
        [0.0000, 0.0000, 0.0000]], device='cuda:0', requires_grad=True)
```

### Versions

Collecting environment information...
PyTorch version: 1.11.0
Is debug build: False
CUDA used to build PyTorch: 11.3
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.3 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.13 (default, Mar 28 2022, 11:38:47)  [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-5.15.0-46-generic-x86_64-with-glibc2.17
Is CUDA available: True
CUDA runtime version: 10.1.243
GPU models and configuration:
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 515.43.04
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] torch==1.11.0
[pip3] torchaudio==0.11.0
[pip3] torchvision==0.12.0
[conda] _pytorch_select           0.1                       cpu_0
[conda] blas                      1.0                         mkl
[conda] cudatoolkit               11.3.1               h2bc3f7f_2
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libmklml                  2019.0.5             h06a4308_0
[conda] mkl                       2020.2                      256
[conda] mkl-service               2.3.0            py38he904b0f_0
[conda] mkl_fft                   1.3.0            py38h54f3939_0
[conda] mkl_random                1.1.1            py38h0573a6f_0
[conda] numpy                     1.19.2           py38h54aff64_0
[conda] numpy-base                1.19.2           py38hfa32c7d_0
[conda] pytorch                   1.11.0          py3.8_cuda11.3_cudnn8.2.0_0    pytorch
[conda] pytorch-mutex             1.0                        cuda    pytorch
[conda] torchaudio                0.11.0               py38_cu113    pytorch
[conda] torchvision               0.12.0               py38_cu113    pytorch

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