RPC - dynamic world size

ItamarWilf: <p>Is it possible to have a dynamic <code>world size</code> when using <code>torch.distributed.rpc</code>?<br>
I want to have a changing number of processes communicating using the TensorPipe backend, without explicitly stating a <code>world size</code>, having each process dynamicaly assigned a rank.</p>

mrshenli: <p>Hey <a class="mention" href="/u/itamarwilf">@ItamarWilf</a>,</p>
<p>Unfortunately, this is not yet possible with the RPC package, but it is in our roadmap.</p>

sagebei: <p>Hi Shen. Is this feature available now? If not, will it be available in the near future?</p>

H-Huang: <p>Hi <a class="mention" href="/u/sagebei">@sagebei</a>, we have a prototype of this feature in the PyTorch 1.12 release which will be available on <strong>June 28</strong> (view releases here: <a href="https://github.com/pytorch/pytorch/releases" class="inline-onebox" rel="noopener nofollow ugc">Releases · pytorch/pytorch · GitHub</a>). Or feel free to pull the nightly PyTorch build.</p>
<p>As part of this feature, an RPC processes <code>init_rpc</code> remains the same; however, when the world_size argument is not specified then this is assumed to be a dynamic group which allows processes to join and leave the group. <code>shutdown()</code> is used to leave the group. We do not currently support dynamic rank allocation. Documentation will be updated for this as well.</p>
<h4>
<a name="current-1" class="anchor" href="#current-1"></a>Current</h4>
<pre><code class="lang-python"># blocking join for all processes
init_rpc("workerN", world_size=N, rank=N)

# blocking shutdown for all processes
shutdown(graceful=True)
# nonblocking shutdown for a single processes
shutdown(graceful=False)
</code></pre>
<h4>
<a name="new-2" class="anchor" href="#new-2"></a>New</h4>
<pre><code class="lang-python"># node join
init_rpc("worker0", rank=0)

# node leave
rpc.shutdown()
</code></pre>
<p>Curious, can you provide a few details on how you will be using this feature and the scenarios / architectures / models you are training? Thanks!</p>

zhanghongyi: <p>Hi <a class="mention" href="/u/h-huang">@H-Huang</a>  <a class="mention" href="/u/mrshenli">@mrshenli</a> ,<br>
I have some questions related to dynamic world size. We are developing a distributed graph learning library for PyTorch <a href="https://github.com/alibaba/graphlearn-for-pytorch" rel="noopener nofollow ugc">graphlearn-for-pytorch</a>. Our current goal is to refactor the library’s RPC functionality and make it dynamic. Sampler processes sample mini-batches from a large graph and serve as background service. Trainer processes can join dynamically and request mini-batches from the sampling processes. This will allow us to run different experiments without reloading the graph.</p>
<p>I have the following questions:</p>
<ol>
<li>Is dynamic world size a stable and long-term supported feature? While I can identify relevant code changes, I couldn’t find any documentation or tutorials.</li>
<li>How are some RPC APIs expected to behave in this dynamic scenario? For example, I tested <code>rpc.api._all_gather</code>, and found that only one rank can receive the response.</li>
</ol>
<pre><code class="lang-auto">import torch.distributed.rpc as rpc
import time
import multiprocessing as mp
import torch
import os

os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '11111'

def worker(rank):
    rpc.init_rpc(
        name=f'worker_{rank}',
        rank=rank,
    )
    print(f'rank {rank} initialized')
    time.sleep(1)
    print(rpc.api._all_gather(torch.cuda.device_count()))
    if rank == 0:
        time.sleep(1)
    rpc.shutdown()
    print(f'rank {rank} exited')

if __name__ == '__main__':
    ranks = [0, 1, 2, 3]
    for rank in ranks:
        process = mp.Process(target=worker, args=(rank, ))
        process.start()
</code></pre>

Wang.Zhang: <p>Hi <a class="mention" href="/u/h-huang">@H-Huang</a> , thank you so much for implementing such feature. I spot this on PyTorch 1.13.1 and hope this feature could help us tremendously.</p>
<p>Our user cases are listed below:</p>
<ol>
<li>adopting spot instance when available</li>
<li>make training task fault-tolerant with dynamic rpc</li>
</ol>
<p>However, there still one issue I found when trying this feature: <code>rpc.shutdown()</code> executes before any rpc request received from other rpc node.</p>
<p>Under the static group setting:</p>
<pre><code class="lang-auto"># blocking join for all processes
init_rpc("workerN", world_size=N, rank=N)

# blocking shutdown for all processes
shutdown(graceful=True)
</code></pre>
<p><code>shutdown</code> in fact does two jobs:</p>
<ol>
<li>blockingly wait permission from other rpc nodes to leave</li>
<li>delete local rrefs, etc.</li>
</ol>
<p>However, in the new style:</p>
<pre><code class="lang-auto"># node join
init_rpc("worker0", rank=0)

# node leave
rpc.shutdown()
</code></pre>
<p>the <code>shutdown</code> will be executing without any blocking, which means without any permission from other rpc node.</p>
<p>Not sure if I’m using the dynamic group feature correctly. I checked the code in 1.13.1 about <code>rpc.shutdown</code> and found <code>rpc_sync(worker.name, _update_group_membership, args=(my_worker_info, [], {}, False))</code> can be executed without any shake-hand with <code>worker</code>.</p>

H-Huang: <p>Hi, thanks for trying out the feature, answer to your questions below:</p>
<blockquote>
<ol>
<li>Is dynamic world size a stable and long-term supported feature? While I can identify relevant code changes, I couldn’t find any documentation or tutorials.</li>
</ol>
</blockquote>
<p>Dynamic world size RPC is a prototype feature and current development is paused. There has been a shift in priority on enhancing Store reliability and robustness (both of which are also needed by RPC). However, the team would be open to review pull requests for RPC.</p>
<blockquote>
<ol start="2">
<li>How are some RPC APIs expected to behave in this dynamic scenario? For example, I tested <code>rpc.api._all_gather</code>, and found that only one rank can receive the response.</li>
</ol>
</blockquote>
<p>rpc_sync and rpc_async should world with dynamic RPC, but things that do not work with CUDA RPC (e.g. RRefs) will also not work with dynamic RPC.</p>
<p>For rpc.api._all_gather that is a private API that is used during RPC initialization. There may be some global state that is sets up which prevents it from being reused again on subsequent calls</p>
<p>We would also accept a public facing version of the API and would review the pull request.</p>

zhanghongyi: <p>Thanks for your response and clarification <a class="mention" href="/u/h-huang">@H-Huang</a> . Appreciate that.</p>

zhanghongyi: <p>Hi <a class="mention" href="/u/h-huang">@H-Huang</a> ,<br>
We are planning to implement our system using the dynamic RPC feature. I would greatly appreciate it if you could answer the following questions:</p>
<ol>
<li>Do rpc_sync and rpc_async work identically to the static ones with thorough tests?</li>
<li>Would bugs related to this feature be given priority so that we could receive assistance when encountering bugs related to it?</li>
<li>Are there any plans to resume the development of this feature in the future?</li>
</ol>
<p>Thank you.</p>