<h2 dir="auto">🐛 Bug</h2>
<p dir="auto">Script shared below runs into RRef leaks and pickler issues when running the DistributedOptimizer</p>
<h2 dir="auto">To Reproduce</h2>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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 = &quot;worker{}&quot;.format(self_rank)
    dst_name = &quot;worker{}&quot;.format(dst_rank)

    rpc.init_model_parallel(
        self_name=self_name,
        self_rank=self_rank,
        worker_name_to_id={&quot;worker0&quot;: 0, &quot;worker1&quot;: 1},
        init_method=&quot;file://{}&quot;.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()"><pre class="notranslate"><code class="notranslate">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()
</code></pre></div>
<p dir="auto">The code above fails with: <a href="https://gist.github.com/pritamdamania87/08cf8de0aa95de03ba1e50da89f96996">https://gist.github.com/pritamdamania87/08cf8de0aa95de03ba1e50da89f96996</a></p>
<p dir="auto">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:</p>
<p dir="auto"><a href="https://gist.github.com/pritamdamania87/e637f5f8a5fede8a8102c584191c28b3">https://gist.github.com/pritamdamania87/e637f5f8a5fede8a8102c584191c28b3</a></p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pietern/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pietern">@pietern</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mrshenli/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mrshenli">@mrshenli</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pritamdamania87/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pritamdamania87">@pritamdamania87</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zhaojuanmao/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zhaojuanmao">@zhaojuanmao</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/satgera/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/satgera">@satgera</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gqchen/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gqchen">@gqchen</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/aazzolini/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/aazzolini">@aazzolini</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/rohan-varma/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/rohan-varma">@rohan-varma</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/xush6528/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/xush6528">@xush6528</a></p>
<h2 dir="auto"><g-emoji class="g-emoji" alias="bug" fallback-src="https://github.githubassets.com/images/icons/emoji/unicode/1f41b.png">🐛</g-emoji> Bug / Enhancement</h2>
<p dir="auto">In distributed training, if <code class="notranslate">rpc.join_rpc()</code> is called when the locally created rrefs are not garbage collected, this can result in leakage errors. See the example script below to reproduce.</p>
<h2 dir="auto">To Reproduce</h2>
<ol dir="auto">
<li>Run the following script</li>
</ol>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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(&quot;worker0&quot;, torch.add, args=(torch.ones(2,2),1)), rpc.remote(&quot;worker0&quot;, 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(&quot;worker{}&quot;.format(rank))
    fn(rank, size)
   #  rpc.join_rpc()

if __name__ == &quot;__main__&quot;:
    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()"><pre class="notranslate"><code class="notranslate">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()
</code></pre></div>
<p dir="auto">Results in the error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="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)"><pre class="notranslate"><code class="notranslate">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)
</code></pre></div>
<h2 dir="auto">Expected behavior</h2>
<p dir="auto">Open for discussion, but I don't think we should have a user error here. The script above can be fixed by running <code class="notranslate">join_rpc()</code> 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 <code class="notranslate">gc</code> to modify this. Should we provide a way for the user to say they are done using the rref?</p>
<p dir="auto">A few ideas are: stop holding the remote ref once <code class="notranslate">to_here</code> is called (this makes sense to me, since the user now has the value locally), or use a context manager that wraps <code class="notranslate">init_model_parallel</code> and <code class="notranslate">join_rpc</code> to garbage collect everything that is needed.</p>
<p dir="auto">cc <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pietern/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pietern">@pietern</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/mrshenli/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/mrshenli">@mrshenli</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/pritamdamania87/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/pritamdamania87">@pritamdamania87</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/zhaojuanmao/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/zhaojuanmao">@zhaojuanmao</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/satgera/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/satgera">@satgera</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/rohan-varma/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/rohan-varma">@rohan-varma</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/gqchen/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/gqchen">@gqchen</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/aazzolini/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/aazzolini">@aazzolini</a> <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/xush6528/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/xush6528">@xush6528</a></p>