Hey @ZiyiZhu, thanks for trying out RPC.

Would you think this is possible and let me know how to do this?


One of my thoughts is that if I could make an RRef to each submodule’s output and pass them among the workers.

Solution I
Yes, this is possible, and using RRef is certainly one proper solution. More specifically, we can let worker 0 serve as a master here. Sth like
# On worker 0
emb_lookup_rref = rpc.remote("worker1", EmbeddingTable.forward, args=(input,))
# note that RNN.forward needs to call to_here() on emb_lookup_rref
rnn_result_rref = rpc.remote("worker2", RNN.forward, args=(emb_lookup_rref,))
# similarly Decoder also needs to call to_here() on rnn_result_rref
decoded_rref = rpc.remote("worker1", Decoder.forward, args=(rnn_result_rref,))
final_result = decoded_rref.to_here()

Above should work. Although it would result in several additional light-weight messages to manage internal RRef reference count, it shouldn’t slow down training. Because rpc.remote is non-blockng, it returns an RRef immediately. It’s very likely that the RNN module is already waiting on to_here() to get the embedding lookup result even before the EmbeddingTable finished processing the request. So that there shouldn’t be noticeable delay on the critical path.
Solution II
An alternative solution is to use nested RPC. You can wrap the EmbeddingTable -> RNN -> Decoder into one module forward function (Say MyModel.forward), and then let worker 0 to run rpc.rpc_sync("worker1", MyModel.forward, args=(input,)). Within MyModel.forward, it can also use rpc/remote to communicate with worker 2, sth like:
class MyModel(nn.Module):
    def forward(self, input):
        lookup_result = self.emb(input)
        # here is directly pass the lookup result instead of RRef to RNN
        rnn_result = rpc.rpc_sync("worker2", RNN.forward, inputs=(lookup_result))
        return self.decoder(rnn_result)