Dear @mrshenli,
I have successfully run the RRN RPC tutorial/example shown in the link below:
https://pytorch.org/tutorials/intermediate/rpc_tutorial.html 5
This is very helpful for me to understand the basics of the RPC framework and has demonstrated it very clear how to do distributed training with two machines (nodes).
However, my question is what if we have 3 or more nodes (workers) and we want to split the model into submodels on each machine/node/worker. How should I use the RPC framework to help to pass the intermediate result of the previous worker to the next worker?
Take the RNN tutorial/model as an example:
In the tutorial, basically it does this in the forward pass:
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)
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)
return decoded, hidden


By using the rpc.remote and rpc_sync, I can have the EmbddingTable and do the forward pass remotely on the worker#1 and get the EmbddingTable’s result back locally.


Then I pass the EmbddingTable’s result to my local RNN (worker#0) and get the corresponding RNN result.


I have another Decoder remotely on worker#1 again, and I push the RNN result to that Decoder and then get the result back by using rpc_sync


However, what if I have three workers, worker#0 (local), worker#1 and worker#2. But this time, I put the RNN model on remote worker#2 and I want to have the communication like below:


From worker#0, I push the input to the EmbddingTable on worker#1. After worker#1 finishes the calculation, it passes the result to the RNN on worker#2.


The RNN on worker#2 calculates and passes the result back to worker#1 for the Decoder.


After the Decoder on worker#1 finishes the computation, I (on worker 0) use rpc_sync or to_here to get the final result back to local.


Would you think this is possible and let me know how to do this? Besides, can the Distributed Autograd and the Distributed Optimizer still be applied in this scenario?
One of my thoughts is that if I could make an RRef to each submodule’s output and pass them among the workers.
Thank you very much in advance for your time and help!
Best,
Ziyi