Issue 1:
🐛 Bug

Script shared below runs into RRef leaks and pickler issues when running the DistributedOptimizer

To Reproduce

import torch
from torch.distributed import rpc
import torch.distributed.autograd as dist_autograd
from torch import optim
from torch.distributed.optim import DistributedOptimizer
from tempfile import NamedTemporaryFile
import multiprocessing as mp

def random_tensor():
return torch.rand((3, 3), requires_grad=True)

def _run_process(self_rank, dst_rank, file_name):
self_name = "worker{}".format(self_rank)
dst_name = "worker{}".format(dst_rank)

rpc.init_model_parallel(
self_name=self_name,
self_rank=self_rank,
worker_name_to_id={"worker0": 0, "worker1": 1},
init_method="file://{}".format(file_name),
)

with dist_autograd.context() as context_id:
# Forward pass.
rref1 = rpc.remote(dst_name, random_tensor)
rref2 = rpc.remote(dst_name, random_tensor)
loss = rref1.to_here() + rref2.to_here()

# Backward pass.
dist_autograd.backward([loss.sum()])

# Optimizer.
try:
dist_optim = DistributedOptimizer(
optim.SGD,
[rref1, rref2],
lr=0.05,
)
except Exception as e:
print (e)
raise e
dist_optim.step()

rpc.join_rpc()

file_name = NamedTemporaryFile().name
processes = []
for i in range(2):
p = mp.Process(target=_run_process, args=(i, (i + 1) % 2, file_name))
p.start()

for p in processes:
p.join()

The code above fails with: https://gist.github.com/pritamdamania87/08cf8de0aa95de03ba1e50da89f96996

It looks like when we create RRefs in the DistributedOptimizer constructor, we don't "wait" for these RRefs. Aside from the issue above, we should ensure we wait for the rrefs so that we discover any errors earlier. But now, if I change the DistributedOptimizer code to wait for rrefs, I run into the following error with the same script above:

https://gist.github.com/pritamdamania87/e637f5f8a5fede8a8102c584191c28b3

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @gqchen @aazzolini @rohan-varma @xush6528

Issue 2:
🐛 Bug / Enhancement

In distributed training, if rpc.join_rpc() is called when the locally created rrefs are not garbage collected, this can result in leakage errors. See the example script below to reproduce.

To Reproduce

Run the following script

import os
import time
from datetime import timedelta
import torch
import torch.distributed as dist
from torch.multiprocessing import Process
import torch.distributed.rpc as rpc

def run(rank, size):
if rank == 1:
rref1, rref2 = rpc.remote("worker0", torch.add, args=(torch.ones(2,2),1)), rpc.remote("worker0", torch.add, args=(torch.zeros(2,2), 3))
x = rref1.to_here() + rref2.to_here()
print(x)
rpc.join_rpc() # comment this and uncomment in init_process to remove error

def init_process(rank, size, fn, backend='gloo'):
os.environ['MASTER_ADDR'] = '127.0.0.1'
os.environ['MASTER_PORT'] = '29509'
dist.init_process_group(backend, rank=rank, world_size=size, timeout=timedelta(seconds=12))
rpc.init_model_parallel("worker{}".format(rank))
fn(rank, size)
# rpc.join_rpc()

if __name__ == "__main__":
size = 2
processes = []
for rank in range(size):
p = Process(target=init_process, args=(rank, size, run))
p.start()
processes.append(p)
for p in processes:
p.join()

Results in the error:

RuntimeError: Leaking RRef GloballyUniqueId(1, 1) with fork Id GloballyUniqueId(1, 0)
Leaking RRef GloballyUniqueId(1, 3) with fork Id GloballyUniqueId(1, 2)
(checkRRefLeaks at ../torch/csrc/distributed/rpc/rref_context.cpp:48)

Expected behavior

Open for discussion, but I don't think we should have a user error here. The script above can be fixed by running join_rpc() after the rrefs are garbage collected (such as in the parent function), but this does not seem robust. In particular, there are no real guarantees (AFAIK) of when exactly the garbage collection will happen and the user can use gc to modify this. Should we provide a way for the user to say they are done using the rref?

A few ideas are: stop holding the remote ref once to_here is called (this makes sense to me, since the user now has the value locally), or use a context manager that wraps init_model_parallel and join_rpc to garbage collect everything that is needed.

cc @pietern @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @xush6528