Does TorchScript support RPC-based models?

Thiago.Crepaldi: <p>I’ve been trying to run <code>torch.jit.script(rpc_model)</code> on the RNN parallel model described at <a href="https://pytorch.org/tutorials/intermediate/rpc_tutorial.html#distributed-rnn-using-distributed-autograd-and-distributed-optimizer" rel="nofollow noopener">Pytorch’s RPC tutorial</a>, but have been hitting issues.</p>
<p>My attempt includes adding <code>script_model = torch.jit.script(model)</code> tight before training loop and comment out the loop, as I dont need to train the model.</p>
<p>The first error I got is related to TorchScript not supportting <code>*args</code> and <code>**kwargs</code> from <code>_remote_method</code> and <code>_call_method</code>. In an attempt to bypass this error, I have changed RNNModel.forward to a simple <code>emb = rpc_sync(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input)</code> (and commented out all other calls to local and remote layers) just to see whether this layer could be converted to torchscript. It failed with the following a similar issue, but not due to any public facing dictionary: <code>torch.jit.frontend.NotSupportedError: keyword-arg expansion is not supported:</code>. Here is the full stack</p>
<p>$ mpirun -np 2 python rnn_jit.py<br>
Traceback (most recent call last):<br>
File “rnn_jit.py”, line 156, in <br>
run_worker()<br>
File “rnn_jit.py”, line 145, in run_worker<br>
_run_trainer()<br>
File “rnn_jit.py”, line 116, in _run_trainer<br>
traced = torch.jit.script(model)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1261, in script<br>
return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 305, in create_script_module<br>
return create_script_module_impl(nn_module, concrete_type, stubs_fn)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 361, in create_script_module_impl<br>
create_methods_from_stubs(concrete_type, stubs)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 279, in create_methods_from_stubs<br>
concrete_type._create_methods(defs, rcbs, defaults)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 568, in try_compile_fn<br>
return torch.jit.script(fn, _rcb=rcb)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1290, in script<br>
fn = torch._C._jit_script_compile(qualified_name, ast, _rcb, get_default_args(obj))<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 568, in try_compile_fn<br>
return torch.jit.script(fn, _rcb=rcb)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1287, in script<br>
ast = get_jit_def(obj)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 173, in get_jit_def<br>
return build_def(ctx, py_ast.body[0], type_line, self_name)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 206, in build_def<br>
build_stmts(ctx, body))<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 129, in build_stmts<br>
stmts = [build_stmt(ctx, s) for s in stmts]<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 129, in <br>
stmts = [build_stmt(ctx, s) for s in stmts]<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 181, in <strong>call</strong><br>
return method(ctx, node)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 363, in build_If<br>
build_stmts(ctx, stmt.body),<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 129, in build_stmts<br>
stmts = [build_stmt(ctx, s) for s in stmts]<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 129, in <br>
stmts = [build_stmt(ctx, s) for s in stmts]<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 181, in <strong>call</strong><br>
return method(ctx, node)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 288, in build_Assign<br>
rhs = build_expr(ctx, stmt.value)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 181, in <strong>call</strong><br>
return method(ctx, node)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/frontend.py”, line 464, in build_Call<br>
raise NotSupportedError(kw_expr.range(), ‘keyword-arg expansion is not supported’)<br>
torch.jit.frontend.NotSupportedError: keyword-arg expansion is not supported:<br>
File “/opt/conda/lib/python3.7/site-packages/torch/distributed/rpc/api.py”, line 476</p>
<pre><code>if qualified_name is not None:
    fut = _invoke_rpc_builtin(dst_worker_info, qualified_name, rf, *args, **kwargs)
                                                                            ~~~~~~ &lt;--- HERE
elif isinstance(func, torch.jit.ScriptFunction):
    fut = _invoke_rpc_torchscript(dst_worker_info.name, func, args, kwargs)
</code></pre>
<p>‘rpc_sync’ is being compiled since it was called from ‘RNNModel.forward’<br>
File “rnn_jit.py”, line 68<br>
# pass input to the remote embedding table and fetch emb tensor back<br>
# emb = _remote_method(EmbeddingTable.forward, self.emb_table_rref, input)<br>
emb = rpc_sync(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input)<br>
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;— HERE<br>
# output, hidden = self.rnn(emb, hidden)<br>
# pass output to the rremote decoder and get the decoded output back</p>
<hr>
<h2>Primary job  terminated normally, but 1 process returned<br>
a non-zero exit code. Per user-direction, the job has been aborted.</h2>
<hr>
<p>mpirun detected that one or more processes exited with non-zero status, thus causing<br>
the job to be terminated. The first process to do so was:</p>
<h2>Process name: [[29685,1],1]<br>
Exit code:    1</h2>

mrshenli: <p>Hey <a class="mention" href="/u/thiago.crepaldi">@Thiago.Crepaldi</a></p>
<p>TorchScript integration with RPC is still experimental, and we are working on closing the gaps. Currently, in v1.5, applications can run TorchScript functions using RPC, e.g., <code>rpc_sync/rpc_async/remote(to, my_script_func, args=(...))</code>. But, within a script function only <code>rpc_async</code> can be called. See the test below:</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/dad552666e1aaff174b549a38fce24b517c53d21/torch/testing/_internal/distributed/rpc/jit/rpc_test.py#L146-L154" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/dad552666e1aaff174b549a38fce24b517c53d21/torch/testing/_internal/distributed/rpc/jit/rpc_test.py#L146-L154" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/dad552666e1aaff174b549a38fce24b517c53d21/torch/testing/_internal/distributed/rpc/jit/rpc_test.py#L146-L154</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="146" style="counter-reset: li-counter 145 ;">
<li>@torch.jit.script</li>
<li>def call_rpc_with_profiling(handle: Tensor, dst_worker_name: str) -&gt; Tensor:</li>
<li>    # Call rpc_async from within ScriptFunction and ensure that we can attach</li>
<li>    # profiling callbacks. Note that handle here is a Tensor representation of</li>
<li>    # RecordFunction.</li>
<li>    fut = rpc.rpc_async(dst_worker_name, one_arg, (torch.tensor(1),))</li>
<li>    torch.ops.profiler._call_end_callbacks_on_jit_fut(handle, fut)</li>
<li>    ret = fut.wait()</li>
<li>    return ret</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>Could you please try if using <code>rpc_async</code> works for you?</p>
<p>The issue below is tracking the progress on adding native <code>RemoteModule</code> support. Please feel free to comment there.</p>
<aside class="onebox githubissue">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/issues/37136" target="_blank" rel="nofollow noopener">github.com/pytorch/pytorch</a>
  </header>
  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewBox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/37136" target="_blank" rel="nofollow noopener">[Design][RFC] RemoteModule API Design</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2020-04-23" data-time="08:03:15" data-timezone="UTC">08:03AM - 23 Apr 20 UTC</span>
      </div>


      <div class="user">
        <a href="https://github.com/xush6528" target="_blank" rel="nofollow noopener">
          <img alt="xush6528" src="https://avatars2.githubusercontent.com/u/7608630?v=4" class="onebox-avatar-inline" width="20" height="20">
          xush6528
        </a>
      </div>
    </div>
  </div>
</div>

<div class="github-row">
  <p class="github-content">[Design][RFC] RemoteModule API Design
FB internal link, [Design] RemoteModule
Goal
Provide a wrapper nn.Module API that serves these purposes
Users should feel like using remote...</p>
</div>

<div class="labels">
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">jit</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">module: rpc</span>
    <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">triaged</span>
</div>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>


Thiago.Crepaldi: <p>Thanks for the quick reply, Shen Li. By using rpc_async, I got similar error:</p>
<p>$ mpirun -np 2 python rnn_jit.py<br>
Traceback (most recent call last):<br>
File “rnn_jit.py”, line 156, in <br>
run_worker()<br>
File “rnn_jit.py”, line 145, in run_worker<br>
_run_trainer()<br>
File “rnn_jit.py”, line 116, in _run_trainer<br>
script_model = torch.jit.script(model)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/<strong>init</strong>.py”, line 1261, in script<br>
return torch.jit._recursive.create_script_module(obj, torch.jit._recursive.infer_methods_to_compile)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 305, in create_script_module<br>
return create_script_module_impl(nn_module, concrete_type, stubs_fn)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 361, in create_script_module_impl<br>
create_methods_from_stubs(concrete_type, stubs)<br>
File “/opt/conda/lib/python3.7/site-packages/torch/jit/_recursive.py”, line 279, in create_methods_from_stubs<br>
concrete_type._create_methods(defs, rcbs, defaults)<br>
RuntimeError:<br>
rpc_async(dst_worker_name, user_callable, args, kwargs)does not support kwargs yet:<br>
File “rnn_jit.py”, line 68<br>
# pass input to the remote embedding table and fetch emb tensor back<br>
# emb = _remote_method(EmbeddingTable.forward, self.emb_table_rref, input)<br>
emb = rpc_async(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input)<br>
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ &lt;— HERE<br>
# output, hidden = self.rnn(emb, hidden)<br>
# pass output to the rremote decoder and get the decoded output back</p>
<p>Line 68 refers to the rpc_async call inside the forward method<br>
emb = rpc_async(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input)</p>

Thiago.Crepaldi: <p>This script is a repro for the issue</p>
<pre><code class="lang-python">import torch
import os
import torch.distributed as dist
from torch.distributed.rpc import rpc_sync
import torch.nn as nn
import torch.nn.functional as F
import torch.distributed.rpc as rpc
from torch import optim #XXX include this
from torch.distributed.optim import DistributedOptimizer #XXX include this
from torch.distributed.rpc import RRef, rpc_async, remote
import torch.distributed.autograd as dist_autograd #XXX include this

def get_local_rank():
    return int(os.environ['OMPI_COMM_WORLD_LOCAL_RANK'])

def _parameter_rrefs(module):
    param_rrefs = []
    for param in module.parameters():
        param_rrefs.append(RRef(param))
    return param_rrefs

class EmbeddingTable(nn.Module):
    r"""
    Encoding layers of the RNNModel
    """
    def __init__(self, ntoken, ninp, dropout):
        super(EmbeddingTable, self).__init__()
        self.drop = nn.Dropout(dropout)
        self.encoder = nn.Embedding(ntoken, ninp).cuda()
        self.encoder.weight.data.uniform_(-0.1, 0.1)

    def forward(self, input):
        return self.drop(self.encoder(input.cuda())).cpu() # XXX: extra ')'


class Decoder(nn.Module):
    def __init__(self, ntoken, nhid, dropout):
        super(Decoder, self).__init__()
        self.drop = nn.Dropout(dropout)
        self.decoder = nn.Linear(nhid, ntoken)
        self.decoder.bias.data.zero_()
        self.decoder.weight.data.uniform_(-0.1, 0.1)

    def forward(self, output):
        return self.decoder(self.drop(output))

class RNNModel(nn.Module):
    def __init__(self, ps, ntoken, ninp, nhid, nlayers, dropout=0.5):
        super(RNNModel, self).__init__()

        # setup embedding table remotely
        self.emb_table_rref = rpc.remote(ps, EmbeddingTable, args=(ntoken, ninp, dropout))
        # setup LSTM locally
        self.rnn = nn.LSTM(ninp, nhid, nlayers, dropout=dropout)
        # setup decoder remotely
        self.decoder_rref = rpc.remote(ps, Decoder, args=(ntoken, nhid, dropout))

    def forward(self, input, hidden):
        # pass input to the remote embedding table and fetch emb tensor back
        # emb = _remote_method(EmbeddingTable.forward, self.emb_table_rref, input) # Original call
        emb = rpc_async(self.emb_table_rref.owner(), EmbeddingTable.forward, args=input) # adapted call

        output, hidden = self.rnn(emb, hidden)

        # pass output to the rremote decoder and get the decoded output back
        # decoded = _remote_method(Decoder.forward, self.decoder_rref, output) # Original call
        decoded = rpc_async(self.decoder_rref.owner(), Decoder.forward, args=output) # adapted call
        return decoded, hidden
    def parameter_rrefs(self):
        remote_params = []
        # get RRefs of embedding table
        # remote_params.extend(_remote_method(_parameter_rrefs, self.emb_table_rref)) # Original call
        remote_params.extend(rpc_async(_parameter_rrefs, self.emb_table_rref))

        # create RRefs for local parameters
        remote_params.extend(_parameter_rrefs(self.rnn))

        # get RRefs of decoder
        # remote_params.extend(_remote_method(_parameter_rrefs, self.decoder_rref)) # Original call
        remote_params.extend(rpc_async(_parameter_rrefs, self.decoder_rref)) # Adapted call
        return remote_params

def _run_trainer():
    ntoken = 10
    ninp = 2
    nhid = 3
    nlayers = 4
    model = RNNModel('ps', ntoken, ninp, nhid, nlayers) # XXX: no rnn.
    script_model = torch.jit.script(model)
    print(script_model)

def run_worker():
    world_size=2
    rank=get_local_rank()
    os.environ['MASTER_ADDR'] = '10.123.134.28'
    os.environ['MASTER_PORT'] = '21234'
    if rank == 1:
        rpc.init_rpc("trainer", rank=rank, world_size=world_size)
        _run_trainer()
    else:
        rpc.init_rpc("ps", rank=rank, world_size=world_size)
        # parameter server do nothing
        pass
    # block until all rpcs finish
    rpc.shutdown()

if __name__=="__main__":
    run_worker()
</code></pre>