ZiyiZhu:

Besides, can the Distributed Autograd and the Distributed Optimizer still be applied in this scenario?


Sorry I missed this. Yes, both of them would still work in this case, as long as you wrap the top-level (not nested) RPCs with distributed autograd context. All RPCs originated from that context will propagate the context id, so that autograd and optimizer on different workers will be able to find the context properly.
For Distributed Optimizer, as long as 1) you provide a correct list of param RRefs to its constructor 2) its step() function is wrapped by the correct dist autograd context, it should work. It does not care where those parameters live.
BTW, in the next release, we are making dist autograd/optimizer functional. They will take the context id as an argument and does not need to be wrapped by a with context statement anymore.