Synchronization issue when combining DPP and RPC - "Parameter marked twice"
Labels: oncall: distributed

### 🐛 Describe the bug

I'm using DDP in conjunction with RPC to create a master-worker setup to train a DQN agent using multiple GPUs in parallel. The master process manages the game and memory buffer and fits one of the partitions of the data, while the workers only fit their respective partitions of the data. I'm training on 2 NVIDIA 3090 GPUs within a single machine.

The minimal code to reproduce the error is shown below. The problem appears to be that the master process triggers a second call to the worker's `_train_on_batch` before the first call is completed. As a result, the parameters of the worker get marked twice, which throws the error. This seems like a race condition since the error isn't reproduced 100% of the time. I added some prints to the code (not shown in the example for clarity purposes) which confirm this is the issue:
```
2023-05-03 20:48:59,315 [INFO] rank=0 | start local_iter=0 | master_iter=0
2023-05-03 20:48:59,330 [INFO] rank=1 | start local_iter=0 | master_iter=0
2023-05-03 20:48:59,935 [INFO] rank=0 | bckwd done local_iter=0 | master_iter=0
2023-05-03 20:48:59,937 [INFO] rank=0 | optim done local_iter=0 | master_iter=0
2023-05-03 20:48:59,938 [INFO] rank=0 | start local_iter=1 | master_iter=1
2023-05-03 20:48:59,939 [INFO] rank=1 | start local_iter=0 | master_iter=1
2023-05-03 20:48:59,959 [INFO] rank=1 | bckwd done local_iter=0 | master_iter=0
2023-05-03 20:48:59,959 [INFO] rank=1 | bckwd done local_iter=0 | master_iter=1
On WorkerInfo(id=1, name=worker_1):
RuntimeError('Expected to mark a variable ready only once. This error is caused [...]
```
In the logs above, `rank` is the process rank (`0` for master, `1` for worker), `local_iter` is the `_train_on_batch` call number (i.e `0` for the first call, `1` for the second call), and `master_iter` is the number of `_train_on_batch` calls the master process has completed (iterations numbers are incremented just before exiting `_train_on_batch`). `backwd done` and `optim done` means the process has returned from the `loss.backward` and `optimizer.step` calls, respectively. We see that the master process has finished it's first training iteration when it signals the worker to start its second iteration, while the latter hasn't finished its first iteration yet.

The error is fixed adding a `dist.barrier` before exiting `_train_on_batch` method, which points towards this being a process synchronization issue. However, the DDP documentation states that 
> Constructor, forward method, and differentiation of the output (or a function of the output of this module) are distributed synchronization points. Take that into account in case different processes might be executing different code.

Thus, shouldn't the processes be synchronized on the `preds = self.model(batch.data.to(self.rank))` call? Or is the `dist.barrier` usage intended in this use case?   

Minimal code to reproduce the error (might not be reproduced 100% of the time - bigger models tend to show more failures):
```python
import os
from dataclasses import dataclass

import torch
import torch.nn as nn
import torch.distributed as dist
import torch.distributed.rpc as rpc
import torch.multiprocessing as mp
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.optim import Optimizer


MASTER_ADDR = "localhost"
MASTER_PORT = "29500"
BACKEND = "nccl"


def _set_master_addr() -> None:
    os.environ["MASTER_ADDR"] = MASTER_ADDR
    os.environ["MASTER_PORT"] = MASTER_PORT


@dataclass
class Batch:
    data: torch.Tensor
    labels: torch.Tensor


class TrainOnBatchMixin:
    def _train_on_batch(self, batch: Batch) -> float:
        preds = self.model(batch.data.to(self.rank))
        self.optimizer.zero_grad()
        batch_loss = self.loss(preds, batch.labels.to(self.rank))
        batch_loss.backward()
        self.optimizer.step()


class Trainer(TrainOnBatchMixin):
    def __init__(
        self,
        model: nn.Module,
        optimizer: Optimizer,
        in_feats: int,
        out_feats: int,
        loss: nn.Module = nn.MSELoss(),
        num_workers: int = 1,
        num_epochs: int = 5,
    ) -> None:
        self.model = model
        self.optimizer = optimizer
        self.in_feats = in_feats
        self.out_feats = out_feats
        self.loss = loss
        self.num_workers = num_workers
        self.num_epochs = num_epochs

        self.rank = 0
        self.world_size = self.num_workers + 1

        self._setup()

    def _setup(self) -> None:
        # spawn workers
        for rank in range(1, self.world_size):
            worker_process = mp.Process(
                target=Worker.init_process,
                kwargs={"rank": 1, "world_size": self.world_size},
            )
            worker_process.start()

        # initialize self
        _set_master_addr()

        # init self RPC
        rpc.init_rpc("master", rank=0, world_size=self.world_size)

        # init self DDP
        dist.init_process_group(rank=0, world_size=self.world_size, backend=BACKEND)

        self.worker_rrefs = [
            rpc.remote(
                f"worker_{rank}",
                Worker,
                kwargs={
                    "rank": rank,
                    "model": self.model,
                    "optimizer": self.optimizer,
                    "loss": self.loss,
                },
            )
            for rank in range(1, self.world_size)
        ]
        self.model = DDP(self.model.to(0), device_ids=[0])

    def train(self) -> None:
        for num_epoch in range(self.num_epochs):
            # create fake data
            data = torch.tensor(
                [[1] * self.in_feats, [12] * self.in_feats], dtype=torch.float32
            )
            labels = torch.tensor(
                [[1] * self.out_feats, [13] * self.out_feats], dtype=torch.float32
            )
            batch = Batch(data, labels)

            # execute workers
            for worker_rref in self.worker_rrefs:
                rpc.rpc_async(
                    worker_rref.owner(),
                    worker_rref.rpc_sync()._train_on_batch,
                    args=(batch,),
                )
            # execute self
            self._train_on_batch(batch)

        rpc.shutdown()


class Worker(TrainOnBatchMixin):
    def __init__(
        self,
        rank: int,
        model: nn.Module,
        optimizer: Optimizer,
        loss: nn.Module,
    ) -> None:
        self.rank = rank
        self.model = DDP(model.to(rank), device_ids=[rank])
        self.optimizer = self._copy_optimizer(optimizer, self.model)
        self.loss = type(loss)()

    @staticmethod
    def _copy_optimizer(optimizer: Optimizer, model: nn.Module) -> Optimizer:
        return type(optimizer)(
            params=[p for p in model.parameters()], **optimizer.defaults
        )

    @staticmethod
    def init_process(
        rank: int,
        world_size: int,
    ) -> None:
        _set_master_addr()

        # init RPC
        rpc.init_rpc(name=f"worker_{rank}", rank=rank, world_size=world_size)

        # init DDP
        dist.init_process_group(rank=rank, world_size=world_size, backend=BACKEND)

        # block until trainer has given shut down signal
        rpc.shutdown()


if __name__ == "__main__":
    in_feats = int(1e5)
    out_feats = 10
    model = nn.Linear(in_feats, out_feats)
    optimizer = torch.optim.SGD(params=model.parameters(), lr=1e-3)
    trainer = Trainer(
        model=model, optimizer=optimizer, in_feats=in_feats, out_feats=out_feats
    )
    trainer.train()

```

### Versions

PyTorch version: 2.0.0+cu117
Is debug build: False
CUDA used to build PyTorch: 11.7
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 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.26.3
Libc version: glibc-2.31

Python version: 3.9.4 (default, Apr 23 2023, 21:06:21)  [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-5.4.0-144-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 10.1.243
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce RTX 3090
GPU 1: NVIDIA GeForce RTX 3090

Nvidia driver version: 495.29.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Byte Order:                      Little Endian
Address sizes:                   39 bits physical, 48 bits virtual
CPU(s):                          16
On-line CPU(s) list:             0-15
Thread(s) per core:              2
Core(s) per socket:              8
Socket(s):                       1
NUMA node(s):                    1
Vendor ID:                       GenuineIntel
CPU family:                      6
Model:                           158
Model name:                      Intel(R) Core(TM) i9-9900K CPU @ 3.60GHz
Stepping:                        12
CPU MHz:                         4868.364
CPU max MHz:                     5000.0000
CPU min MHz:                     800.0000
BogoMIPS:                        7200.00
Virtualization:                  VT-x
L1d cache:                       256 KiB
L1i cache:                       256 KiB
L2 cache:                        2 MiB
L3 cache:                        16 MiB
NUMA node0 CPU(s):               0-15
Vulnerability Itlb multihit:     KVM: Mitigation: Split huge pages
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:          Mitigation; IBRS
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Mitigation; Microcode
Vulnerability Tsx async abort:   Mitigation; Clear CPU buffers; SMT vulnerable
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx smx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx rdseed adx smap clflushopt intel_pt xsaveopt xsavec xgetbv1 xsaves dtherm ida arat pln pts hwp hwp_notify hwp_act_window hwp_epp md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] flake8==4.0.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.3
[pip3] torch==2.0.0
[pip3] triton==2.0.0
[conda] Could not collect

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