Repository: pytorch
Issue: DDP with cuda rpc failed with DistributedOptimizer Adagrad
Labels: oncall: distributed,module: optimizer

## 🐛 Bug
I follow tutorial and try to combine DDP with cuda rpc.
https://pytorch.org/tutorials/advanced/rpc_ddp_tutorial.html
https://pytorch.org/tutorials/recipes/cuda_rpc.html

Works well with `SGD`, `Adam` and fail when using `Adagrad`.

## To Reproduce

Steps to reproduce the behavior:

1. Run following code snippets. (Success when use SGD or Adam)

```python
import random

import torch
import torch.distributed as dist
import torch.distributed.autograd as dist_autograd
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
import torch.optim as optim
from torch import nn
from torch.distributed.nn import RemoteModule
from torch.distributed.optim import DistributedOptimizer
from torch.distributed.rpc import RRef
from torch.distributed.rpc import TensorPipeRpcBackendOptions
from torch.nn.parallel import DistributedDataParallel as DDP

NUM_EMBEDDINGS = 100
EMBEDDING_DIM = 2
batch_size = 100


class HybridModel(torch.nn.Module):
  r"""
    The model consists of a sparse part and a dense part.
    1) The dense part is an nn.Linear module that is replicated across all trainers using DistributedDataParallel.
    2) The sparse part is a Remote Module that holds an nn.EmbeddingBag on the parameter server.
    This remote model can get a Remote Reference to the embedding table on the parameter server.
    """

  def __init__(self, remote_emb_module, device):
    print(f"Init HybridModel on device {device}")
    super(HybridModel, self).__init__()
    self.remote_emb_module = remote_emb_module
    self.fc = DDP(torch.nn.Linear(EMBEDDING_DIM, 8).cuda(device),
                  device_ids=[device])
    self.device = device

  def forward(self, indices):
    emb_lookup1 = self.remote_emb_module[0].forward(indices[:, 0])
    emb_lookup2 = self.remote_emb_module[1].forward(indices[:, 1])
    if type(emb_lookup1) in (tuple, list):
      emb_lookup1 = torch.cat(emb_lookup1, dim=0).reshape(-1, EMBEDDING_DIM)
    if type(emb_lookup2) in (tuple, list):
      emb_lookup2 = torch.cat(emb_lookup2, dim=0).reshape(-1, EMBEDDING_DIM)

    print(f"curr rank: {torch.distributed.get_rank()}, "
          f"self.device: {self.device}, "
          f"indices[:, 0] shape: {indices[:, 0].shape},\n"
          f"emb_lookup1 shape: {emb_lookup1.shape},\n"
          f"emb_lookup2 shape: {emb_lookup2.shape},\n")
    return self.fc(
        emb_lookup1.cuda(self.device) + emb_lookup2.cuda(self.device))


def _run_trainer(remote_emb_module, rank):
  r"""
    Each trainer runs a forward pass which involves an embedding lookup on the
    parameter server and running nn.Linear locally. During the backward pass,
    DDP is responsible for aggregating the gradients for the dense part
    (nn.Linear) and distributed autograd ensures gradients updates are
    propagated to the parameter server.
    """
  print(f"Setup the model on {rank}")
  model = HybridModel(remote_emb_module=remote_emb_module, device=rank)

  model_parameter_rrefs = []
  for rm in model.remote_emb_module:
    model_parameter_rrefs += rm.remote_parameters()

  for param in model.fc.parameters():
    model_parameter_rrefs.append(RRef(param))

  opt = DistributedOptimizer(
      optim.Adagrad,
      model_parameter_rrefs,
      lr=0.05,
  )

  criterion = torch.nn.CrossEntropyLoss()

  def get_next_batch(rank):
    for _ in range(10):
      indices = torch.LongTensor(100, 2).random_(0, NUM_EMBEDDINGS)
      target = torch.LongTensor(100).random_(8).cuda(rank)
      yield indices, target

  for epoch in range(100):
    for target_x, target_y in get_next_batch(rank):
      with dist_autograd.context() as context_id:
        y_pred = model(target_x)
        loss = criterion(y_pred, target_y)
        dist_autograd.backward(context_id, [loss])
        opt.step(context_id)

    print("Training done for epoch {}".format(epoch))


def run_worker(rank, world_size):
  r"""
    A wrapper function that initializes RPC, calls the function, and shuts down
    RPC.
    """

  rpc_backend_options = TensorPipeRpcBackendOptions()
  rpc_backend_options.init_method = "tcp://localhost:29501"

  # Rank 2 is master, 3 is ps and 0 and 1 are trainers.
  if rank == 2:
    rpc.init_rpc(
        "master",
        rank=rank,
        world_size=world_size,
        rpc_backend_options=rpc_backend_options,
    )

    remote_emb_module = [
        RemoteModule(
            "trainer0/cuda:0",
            nn.Embedding,
            args=(NUM_EMBEDDINGS, EMBEDDING_DIM),
            kwargs={},
        ),
        RemoteModule(
            "trainer1/cuda:1",
            nn.Embedding,
            args=(NUM_EMBEDDINGS, EMBEDDING_DIM),
            kwargs={},
        ),
    ]

    # Run the training loop on trainers.
    futs = []
    for trainer_rank in [0, 1]:
      trainer_name = "trainer{}".format(trainer_rank)
      fut = rpc.rpc_async(trainer_name,
                          _run_trainer,
                          args=(remote_emb_module, trainer_rank))
      futs.append(fut)

    # Wait for all training to finish.
    for fut in futs:
      fut.wait()
  elif rank <= 1:
    # Initialize process group for Distributed DataParallel on trainers.
    dist.init_process_group(backend="gloo",
                            rank=rank,
                            world_size=2,
                            init_method="tcp://localhost:29500")

    # Initialize RPC.
    trainer_name = "trainer{}".format(rank)
    worker_rpc_backend_options = TensorPipeRpcBackendOptions()
    worker_rpc_backend_options.init_method = "tcp://localhost:29501"
    for remote_rank in range(2):
      if remote_rank != rank:
        worker_rpc_backend_options.set_device_map(f"trainer{remote_rank}",
                                                  {rank: remote_rank})
    rpc.init_rpc(
        trainer_name,
        rank=rank,
        world_size=world_size,
        rpc_backend_options=worker_rpc_backend_options,
    )

  # Trainer just waits for RPCs from master.

  # block until all rpcs finish
  rpc.shutdown()


if __name__ == "__main__":
  # 2 trainers, 1 parameter server, 1 master.
  world_size = 3
  mp.spawn(run_worker, args=(world_size,), nprocs=world_size, join=True)

```

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

PyTorch version: 1.9.0.dev20210601+cu111
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: CentOS Linux release 7.9.2009 (Core) (x86_64)
GCC version: (GCC) 8.3.1 20190311 (Red Hat 8.3.1-3)
Clang version: Could not collect
CMake version: version 3.16.5
Libc version: glibc-2.2.5

Python version: 3.7 (64-bit runtime)
Python platform: Linux-3.10.0-1160.15.2.el7.x86_64-x86_64-with-centos-7.9.2009-Core
Is CUDA available: True
CUDA runtime version: 11.1.74
GPU models and configuration:
GPU 0: Tesla T4
GPU 1: Tesla T4
GPU 2: Tesla T4
GPU 3: Tesla T4

Nvidia driver version: 460.32.03
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.5
[pip3] torch==1.9.0.dev20210601+cu111
[pip3] torchaudio==0.9.0.dev20210601
[pip3] torchvision==0.11.0.dev20210601+cu111


## Error message:
```
terminate called after throwing an instance of 'c10::Error'
  what():  device_index >= 0 && device_index < num_gpusINTERNAL ASSERT FAILED at "/pytorch/c10/cuda/CUDAStream.cpp":254, please report a bug to PyTorch.
Exception raised from check_gpu at /pytorch/c10/cuda/CUDAStream.cpp:254 (most recent call first):
frame #0: c10::Error::Error(c10::SourceLocation, std::string) + 0x42 (0x7f563cd26412 in /data0/fumi/.pyenv/versions/venv374/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #1: c10::detail::torchCheckFail(char const*, char const*, unsigned int, char const*) + 0x5f (0x7f563cd22e7f in /data0/fumi/.pyenv/versions/venv374/lib/python3.7/site-packages/torch/lib/libc10.so)
frame #2: <unknown function> + 0x64bf (0x7f56db43c4bf in /data0/fumi/.pyenv/versions/venv374/lib/python3.7/site-packages/torch/lib/libc10_cuda.so)
frame #3: <unknown function> + 0x8fa2 (0x7f56db43efa2 in /data0/fumi/.pyenv/versions/venv374/lib/python3.7/site-packages/torch/lib/libc10_cuda.so)
frame #4: <unknown function> + 0xfcf29c (0x7f563dd4c29c in /data0/fumi/.pyenv/versions/venv374/lib/python3.7/site-packages/torch/lib/libtorch_cpu.so)
```


cc @vincentqb @iramazanli @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @agolynski @SciPioneer @H-Huang @mrzzd @cbalioglu @gcramer23